Source code for aidputils.agents.toolkit.tool_helper
from pydantic import create_model, Field
from aidputils.agents.auth.util import auth_utils
from aidputils.agents.tools import utils as aidp_tool_utils
from aidputils.agents.toolkit import tool_common
from langchain_core.tools import StructuredTool
from typing import Optional, Callable, Any, Awaitable
import mcp.types as mcp_types
import logging
logger = logging.getLogger('tool_helper')
tool_conf = None
# global tool registry for use during langgraph agent creation.
tool_registry = {}
def _error(message: str, code: int, data: dict | None = None) -> dict:
return tool_common.error(message, code, data)
def _as_str(v: Any) -> str:
return tool_common.as_str(v)
def _sanitize_params(params: dict) -> dict:
return tool_common.sanitize_params(params)
[docs]
def create_tool_function(tool_name: str):
"""Create a LangGraph-compatible wrapper that invokes a registered toolkit tool.
The returned callable injects the selected ``tool_name`` into the runtime parameters
and delegates execution to ``aidputils.agents.tools.utils.call_tool_by_name``.
This allows the same registry-backed tool configuration to be reused by both the
toolkit runtime and generated StructuredTool objects.
"""
def specific_tool(**kwargs) -> str:
params = dict(kwargs)
params.update({'tool_name': tool_name})
return aidp_tool_utils.call_tool_by_name(tool_name = tool_name, tool_registry=tool_registry, runtime_params=params)
return specific_tool
[docs]
def derive_primitive_type(type_str):
"""Map a tool-spec type string to a Python runtime type.
This helper normalizes loose type spellings coming from tool metadata so dynamic
Pydantic schemas can be constructed consistently.
Supports:
- Primitives: str/string, int/integer, float/double/number/numeric, bool/boolean, bytes, bytearray
- Nulls: none, null
- Collections: list/array/sequence/iterable, dict/dictionary/map/mapping, set, frozenset, tuple
- Basic generic-like patterns: "list[int]", "set[str]", "tuple[int, str]" or "list of int"
In these cases, we return the base container type (list/set/tuple/dict).
"""
if not isinstance(type_str, str):
return None
t = type_str.strip().lower()
if not t:
return None
# Extract base for simple generic-like notations:
# - list[int], set[str], tuple[int, str], dict[str, int]
# - "list of int", "set of str", etc.
if '[' in t:
t = t.split('[', 1)[0].strip()
elif ' of ' in t:
t = t.split(' of ', 1)[0].strip()
# Normalize common synonyms to canonical tokens
synonym_map = {
'string': 'str',
'integer': 'int',
'double': 'float',
'number': 'float',
'numeric': 'float',
'boolean': 'bool',
'dictionary': 'dict',
'map': 'dict',
'mapping': 'dict',
'array': 'list',
'sequence': 'list',
'iterable': 'list',
}
canonical = synonym_map.get(t, t)
type_map = {
'str': str,
'int': int,
'float': float,
'bool': bool,
'bytes': bytes,
'bytearray': bytearray,
'list': list,
'dict': dict,
'set': set,
'frozenset': frozenset,
'tuple': tuple,
'none': type(None),
'null': type(None),
}
return type_map.get(canonical, None)
[docs]
def create_langgraph_tool(tool):
"""Convert a toolkit tool configuration into a ``StructuredTool`` instance.
The resulting tool uses the global ``tool_registry`` so downstream execution can look
up the original tool metadata and invoke the correct implementation at runtime.
"""
name = tool['name']
description = tool['description']
print(f"Adding tool {name}!")
s_tool = StructuredTool.from_function(
func=create_tool_function(tool_name=name),
# func=tool_method_registry[name],
name=name,
description=description,
args_schema=create_tool_schema(tool),
infer_schema=False
)
# Add the tool config to the tool registry
global tool_registry
tool['class'] = tool['tool_class'] # The tool utils expect 'class' in the config.
tool_registry[name] = tool
return s_tool
[docs]
def create_langgraph_custom_tool(tool):
"""Create a ``StructuredTool`` backed by a user-provided Python callable.
This path is intended for custom tools that do not rely on the standard toolkit
registry invocation flow but should still appear as first-class LangChain tools.
"""
name = tool['name']
description = tool['description']
print(f"Adding tool {name}!")
s_tool = StructuredTool.from_function(
func=tool['conf']['func'],
name=name,
description=description,
args_schema=create_tool_schema(tool),
infer_schema=False
)
# Add the tool config to the tool registry
global tool_registry
tool['class'] = tool['tool_class'] # The tool utils expect 'class' in the config.
tool_registry[name] = tool
return s_tool
[docs]
def create_tool_schema(tool):
"""Build a dynamic Pydantic input model from a toolkit tool parameter definition."""
fields = {}
for param in tool['params']:
name = param['name']
description = param.get('description', '')
fields.update({name: (derive_primitive_type(param['type']), Field(description=description))})
args_schema = create_model("DynamicInput", **fields)
return args_schema
def _py_type_from_jsonschema_type(type_name: Optional[str]):
"""
Map JSON Schema 'type' to a Python type for Pydantic model generation.
"""
if not isinstance(type_name, str):
return Any
t = type_name.lower()
return {
"string": str,
"integer": int,
"number": float,
"boolean": bool,
"array": list,
"object": dict,
"null": type(None),
}.get(t, Any)
[docs]
def build_args_schema_from_json_schema(schema: Optional[dict], model_name: str = "MCPToolInput"):
"""
Convert a JSON Schema (as provided by mcp.types.Tool.inputSchema) into a Pydantic model class.
- Supports basic 'properties' and 'required' handling
- Maps primitive types and falls back to Any for unknowns
"""
if not isinstance(schema, dict):
return create_model(model_name)
props = schema.get("properties") or {}
required = set(schema.get("required") or [])
fields = {}
for name, spec in props.items():
if not isinstance(spec, dict):
continue
py_type = _py_type_from_jsonschema_type(spec.get("type"))
desc = spec.get("description", "")
default = spec.get("default", None)
if name in required:
fields[name] = (py_type, Field(..., description=desc))
else:
# Optional parameter: default to provided default or None
fields[name] = (py_type, Field(default if default is not None else None, description=desc))
return create_model(model_name, **fields)
[docs]
def mcp_tool_to_structured_tool(mcp_tool: "mcp_types.Tool", coroutine: Optional[Callable[..., Awaitable[Any]]] = None,) -> StructuredTool:
"""
Convert an mcp.types.Tool into a LangChain StructuredTool.
- Name/description taken from the MCP tool
- Args schema built from MCP tool's inputSchema JSON schema
- 'func' is the callable to execute when the tool is invoked; if not provided,
a default invoker returns a JSON-serializable envelope with the tool name and arguments.
Example usage:
st = mcp_tool_to_structured_tool(mcp_tool, func=lambda **kw: client.call_tool(...))
"""
name = getattr(mcp_tool, "name", None) or "mcp_tool"
description = getattr(mcp_tool, "description", None) or ""
input_schema = getattr(mcp_tool, "inputSchema", None)
args_schema = build_args_schema_from_json_schema(input_schema, model_name=name)
def _default_invoker(**kwargs):
# Return a simple envelope; caller can intercept and perform the actual call
return {"tool_name": name, "arguments": kwargs}
return StructuredTool.from_function(
coroutine=coroutine or _default_invoker,
name=name,
description=description,
args_schema=args_schema,
infer_schema=False,
)
[docs]
def build_mcp_invoker_factory(client, timeout_secs: Optional[float] = None, runtime_params: dict = None):
"""Build an invoker factory for remotely hosted MCP tools.
Each generated coroutine merges configured default arguments, resolves auth context,
and routes the request through the provided MCP client.
Defaults merging (scoped by server_name to avoid cross-server conflicts):
- allowed_tools entries may include:
- name/toolName: specific tool name
- argOverrides/arg_overrides: dict of default argument values
- server_name/serverName (optional): to explicitly scope the entry
- Merge order per invocation: per-tool defaults (for this server) -> user kwargs (user wins)
"""
# Precompute default arguments from allowed_tools, partitioned by server_name
per_server_per_name_defaults: dict[str, dict[str, dict]] = {}
server_name = (runtime_params or {}).get("server_name") if isinstance(runtime_params, dict) else None
allowed_tools = (runtime_params or {}).get("allowed_tools") if isinstance(runtime_params, dict) else None
for item in allowed_tools:
name = item.get("name")
defs = item.get("argOverrides") or item.get("arg_overrides") or {}
if not isinstance(defs, dict):
continue
per_server_per_name_defaults.setdefault(server_name, {})
per_server_per_name_defaults[server_name][name] = {
**(per_server_per_name_defaults[server_name].get(name, {})),
**defs,
}
def _merge_defaults(tool_name: str, provided: dict) -> dict:
# apply per-tool defaults for this server; user kwargs take precedence
base: dict = {}
server_tool_defs = (per_server_per_name_defaults.get(server_name) or {}).get(tool_name) or {}
if server_tool_defs:
base.update(server_tool_defs)
# finally user-provided args
provided.update(base or {})
return provided
def invoker_factory(tool_name: str):
async def _invoke(**arguments: dict[str, Any]):
context = auth_utils.get_auth_context()
session_id = context.session_id
merged = _merge_defaults(tool_name, arguments)
# disabling session caching
use_cached = False
function_name = "mcp_tool"
from aidputils.agents.tools.tool_traces import TraceUtility
trace_util = TraceUtility("MCP_TOOL_CALL")
span_attributes = {"function": function_name}
if tool_name is not None:
span_attributes.update({"external_tool": tool_name})
span_attributes.update({"server": server_name})
try:
# Trace in parent using start_as_current_span; metrics handled in this parent as well
with trace_util.start_as_current_span(function_name, attributes=span_attributes) as span:
return await client.call_tool(server_name, tool_name, merged, 10, use_cached=use_cached, session_id=session_id)
except Exception as ex:
logger.exception("Failed to invoke tool %s", tool_name)
raise ex
return _invoke
return invoker_factory
[docs]
def build_args_schema_from_json_schema_with_overrides(
schema: Optional[dict],
overrides: Optional[dict],
model_name: str = "MCPToolInput",
):
"""
Build a Pydantic model from a JSON Schema, applying argOverrides:
- For every attribute present in overrides:
* remove it from input schema 'properties'
* remove it from 'required' (so it is not requested from the LLM)
- Do NOT treat overrides as defaults in the schema.
- Then build a Pydantic model from the pruned schema.
"""
if not isinstance(schema, dict):
return create_model(model_name)
# Normalize
props = schema.get("properties") or {}
required = schema.get("required") or []
override_keys = set(overrides.keys()) if isinstance(overrides, dict) else set()
# Prune properties and required for all override keys
if override_keys:
pruned_props = {k: v for k, v in props.items() if k not in override_keys}
# normalize required to list
if isinstance(required, list):
pruned_required = [r for r in required if r not in override_keys]
else:
try:
pruned_required = [r for r in list(required) if r not in override_keys]
except Exception:
pruned_required = []
else:
pruned_props = dict(props)
pruned_required = list(required) if isinstance(required, list) else (list(required) if required is not None else [])
pruned_schema = dict(schema)
pruned_schema["properties"] = pruned_props
pruned_schema["required"] = pruned_required
# Delegate actual Pydantic model construction to the base builder
return build_args_schema_from_json_schema(pruned_schema, model_name=model_name)
[docs]
def build_structured_tools_from_allowed_mcp_tools(
allowed_tools: list[dict],
server_name: Optional[str] = None,
endpoint: Optional[str] = None,
transport: str = "streamable_http",
auth: Optional[dict] = None,
headers: Optional[dict] = None,
invoker_factory: Optional[Callable[[str], Callable[..., Awaitable[Any]]]] = None,
) -> list[StructuredTool]:
"""
Given an array of AllowedToolDetails (from DP spec), return a list of StructuredTool.
AllowedToolDetails schema (relevant parts):
- instruction: optional custom instruction to override description
- argOverrides: map[string]string default values for tool params
- tool: McpToolObject { name, description, inputSchema }
Behavior:
- Build args_schema from tool.inputSchema and apply argOverrides:
* required fields present in argOverrides are made optional with provided defaults
* non-required fields prefer override as default
- Description is overridden by 'instruction' when provided, else tool.description
- If invoker_factory not provided and server_name/endpoint are provided, this function will:
* Construct an MCP client internally (not stored by caller) and
* Build an invoker_factory (via build_mcp_invoker_factory) so tool coroutines call client.call_tool.
No client storage/cleanup is required by the caller.
- If invoker_factory is provided explicitly, it will be used directly.
Returns:
list[StructuredTool]
"""
results: list[StructuredTool] = []
if not isinstance(allowed_tools, list):
return results
# If no invoker_factory supplied, and server/endpoint given, build one internally.
if invoker_factory is None and server_name and endpoint:
try:
from aidputils.agents.tools.mcp.mcp_service import get_mcp_client as _get_mcp_client
# Build allowed defaults for invoker factory
runtime_allowed = []
for it in allowed_tools:
try:
t = (it or {}).get("tool") or {}
nm = t.get("name")
if not nm:
continue
defs = (it or {}).get("argOverrides") or {}
runtime_allowed.append({"name": nm, "argOverrides": defs})
except Exception:
continue
client = _get_mcp_client(
server_name=server_name,
server_url=endpoint,
auth=auth,
transport=transport,
custom_headers=headers,
)
invoker_factory = build_mcp_invoker_factory(
client,
runtime_params={"server_name": server_name, "allowed_tools": runtime_allowed},
)
except Exception:
logger.exception("Failed to construct MCP client/invoker_factory; falling back to default stub invokers")
for item in allowed_tools:
try:
if not isinstance(item, dict):
continue
tool_obj = item.get("tool") or {}
if not isinstance(tool_obj, dict):
continue
name = tool_obj.get("name")
if not name or not isinstance(name, str):
continue
description = item.get("instruction") or tool_obj.get("description") or ""
input_schema = tool_obj.get("inputSchema") or {}
overrides = item.get("argOverrides") or {}
# Build args schema with overrides applied
args_schema = build_args_schema_from_json_schema_with_overrides(
input_schema, overrides, model_name=f"{name}_Input"
)
def _default_invoker(**kwargs):
return {"tool_name": name, "arguments": kwargs}
coroutine = None
if callable(invoker_factory):
try:
coroutine = invoker_factory(name)
except Exception:
coroutine = None
st = StructuredTool.from_function(
coroutine=coroutine or _default_invoker,
name=name,
description=description,
args_schema=args_schema,
infer_schema=False,
)
results.append(st)
except Exception:
# Skip malformed entries to avoid breaking the caller
logger.exception("Failed to convert AllowedToolDetails to StructuredTool for entry: %s", str(item))
continue
return results