Skip to content

Commit df075db

Browse files
PCfVWPCfVW
committed
feat: add Progressive Tool Discovery pattern tools (arango_search_tools, arango_list_tools_by_category)
1 parent 2e0de95 commit df075db

File tree

3 files changed

+259
-2
lines changed

3 files changed

+259
-2
lines changed

‎mcp_arangodb_async/handlers.py‎

Lines changed: 177 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@
8080
import json
8181
import logging
8282
from contextlib import contextmanager
83+
from datetime import datetime
8384
from jsonschema import Draft7Validator, ValidationError as JSONSchemaValidationError
8485
from arango.database import StandardDatabase
8586
from arango.exceptions import ArangoError
@@ -93,7 +94,7 @@
9394
validate_graph_integrity,
9495
calculate_graph_statistics,
9596
)
96-
from .tool_registry import register_tool
97+
from .tool_registry import register_tool, TOOL_REGISTRY
9798
from .tools import (
9899
ARANGO_QUERY,
99100
ARANGO_LIST_COLLECTIONS,
@@ -129,6 +130,17 @@
129130
ARANGO_VALIDATE_GRAPH_INTEGRITY,
130131
ARANGO_GRAPH_STATISTICS,
131132
ARANGO_DATABASE_STATUS,
133+
# Pattern 1: Progressive Tool Discovery
134+
ARANGO_SEARCH_TOOLS,
135+
ARANGO_LIST_TOOLS_BY_CATEGORY,
136+
# Pattern 2: Context Switching
137+
ARANGO_SWITCH_CONTEXT,
138+
ARANGO_GET_ACTIVE_CONTEXT,
139+
ARANGO_LIST_CONTEXTS,
140+
# Pattern 3: Tool Unloading
141+
ARANGO_ADVANCE_WORKFLOW_STAGE,
142+
ARANGO_GET_TOOL_USAGE_STATS,
143+
ARANGO_UNLOAD_TOOLS,
132144
)
133145

134146
# Configure logger for handlers
@@ -238,6 +250,17 @@ def safe_cursor(cursor):
238250
ValidateGraphIntegrityArgs,
239251
GraphStatisticsArgs,
240252
ArangoDatabaseStatusArgs,
253+
# Pattern 1: Progressive Tool Discovery
254+
SearchToolsArgs,
255+
ListToolsByCategoryArgs,
256+
# Pattern 2: Context Switching
257+
SwitchContextArgs,
258+
GetActiveContextArgs,
259+
ListContextsArgs,
260+
# Pattern 3: Tool Unloading
261+
AdvanceWorkflowStageArgs,
262+
GetToolUsageStatsArgs,
263+
UnloadToolsArgs,
241264
)
242265

243266

@@ -1791,8 +1814,160 @@ def handle_arango_database_status(
17911814
}
17921815

17931816

1817+
# ============================================================================
1818+
# MCP Design Pattern Tools
1819+
# ============================================================================
1820+
1821+
# Tool category mappings for Progressive Tool Discovery
1822+
TOOL_CATEGORIES = {
1823+
"core_data": [
1824+
ARANGO_QUERY, ARANGO_LIST_COLLECTIONS, ARANGO_INSERT,
1825+
ARANGO_UPDATE, ARANGO_REMOVE, ARANGO_CREATE_COLLECTION, ARANGO_BACKUP
1826+
],
1827+
"indexing": [
1828+
ARANGO_LIST_INDEXES, ARANGO_CREATE_INDEX, ARANGO_DELETE_INDEX, ARANGO_EXPLAIN_QUERY
1829+
],
1830+
"validation": [
1831+
ARANGO_VALIDATE_REFERENCES, ARANGO_INSERT_WITH_VALIDATION,
1832+
ARANGO_BULK_INSERT, ARANGO_BULK_UPDATE
1833+
],
1834+
"schema": [ARANGO_CREATE_SCHEMA, ARANGO_VALIDATE_DOCUMENT],
1835+
"query": [ARANGO_QUERY_BUILDER, ARANGO_QUERY_PROFILE],
1836+
"graph_basic": [
1837+
ARANGO_CREATE_GRAPH, ARANGO_LIST_GRAPHS, ARANGO_ADD_VERTEX_COLLECTION,
1838+
ARANGO_ADD_EDGE_DEFINITION, ARANGO_ADD_EDGE, ARANGO_TRAVERSE, ARANGO_SHORTEST_PATH
1839+
],
1840+
"graph_advanced": [
1841+
ARANGO_BACKUP_GRAPH, ARANGO_RESTORE_GRAPH, ARANGO_BACKUP_NAMED_GRAPHS,
1842+
ARANGO_VALIDATE_GRAPH_INTEGRITY, ARANGO_GRAPH_STATISTICS
1843+
],
1844+
"aliases": [ARANGO_GRAPH_TRAVERSAL, ARANGO_ADD_VERTEX],
1845+
"health": [ARANGO_DATABASE_STATUS]
1846+
}
1847+
1848+
1849+
# Pattern 1: Progressive Tool Discovery
1850+
@handle_errors
1851+
@register_tool(
1852+
name=ARANGO_SEARCH_TOOLS,
1853+
description="Search for MCP tools by keywords and categories. Enables progressive tool discovery by returning only relevant tools instead of loading all 34 tools upfront.",
1854+
model=SearchToolsArgs,
1855+
)
1856+
def handle_search_tools(
1857+
db: StandardDatabase, args: Dict[str, Any]
1858+
) -> Dict[str, Any]:
1859+
"""Search for tools matching keywords and optional category filters.
1860+
1861+
This tool enables Progressive Tool Discovery pattern by allowing AI agents
1862+
to search for relevant tools on-demand rather than loading all tools upfront.
1863+
1864+
Args:
1865+
db: ArangoDB database instance (not used, but required for handler signature)
1866+
args: Validated SearchToolsArgs containing keywords, categories, and detail_level
1867+
1868+
Returns:
1869+
Dictionary with matching tools and search metadata
1870+
"""
1871+
keywords = [kw.lower() for kw in args["keywords"]]
1872+
categories_filter = args.get("categories")
1873+
detail_level = args.get("detail_level", "summary")
1874+
1875+
# Build list of tools to search
1876+
tools_to_search = []
1877+
if categories_filter:
1878+
for cat in categories_filter:
1879+
if cat in TOOL_CATEGORIES:
1880+
tools_to_search.extend(TOOL_CATEGORIES[cat])
1881+
else:
1882+
# Search all tools
1883+
tools_to_search = list(TOOL_REGISTRY.keys())
1884+
1885+
# Search for matching tools
1886+
matches = []
1887+
for tool_name in tools_to_search:
1888+
if tool_name not in TOOL_REGISTRY:
1889+
continue
1890+
1891+
tool_reg = TOOL_REGISTRY[tool_name]
1892+
search_text = f"{tool_reg.name} {tool_reg.description}".lower()
1893+
1894+
# Check if any keyword matches
1895+
if any(keyword in search_text for keyword in keywords):
1896+
if detail_level == "name":
1897+
matches.append({"name": tool_reg.name})
1898+
elif detail_level == "summary":
1899+
matches.append({
1900+
"name": tool_reg.name,
1901+
"description": tool_reg.description
1902+
})
1903+
else: # full
1904+
matches.append({
1905+
"name": tool_reg.name,
1906+
"description": tool_reg.description,
1907+
"inputSchema": tool_reg.model.model_json_schema()
1908+
})
1909+
1910+
return {
1911+
"matches": matches,
1912+
"total_matches": len(matches),
1913+
"keywords": args["keywords"],
1914+
"categories_searched": categories_filter or "all",
1915+
"detail_level": detail_level
1916+
}
1917+
1918+
1919+
@handle_errors
1920+
@register_tool(
1921+
name=ARANGO_LIST_TOOLS_BY_CATEGORY,
1922+
description="List all MCP tools organized by category. Useful for understanding tool organization and selecting workflow-specific tool sets.",
1923+
model=ListToolsByCategoryArgs,
1924+
)
1925+
def handle_list_tools_by_category(
1926+
db: StandardDatabase, args: Dict[str, Any]
1927+
) -> Dict[str, Any]:
1928+
"""List tools organized by category.
1929+
1930+
Args:
1931+
db: ArangoDB database instance (not used, but required for handler signature)
1932+
args: Validated ListToolsByCategoryArgs with optional category filter
1933+
1934+
Returns:
1935+
Dictionary with tools organized by category
1936+
"""
1937+
category_filter = args.get("category")
1938+
1939+
if category_filter:
1940+
# Return single category
1941+
if category_filter not in TOOL_CATEGORIES:
1942+
return {
1943+
"error": f"Unknown category: {category_filter}",
1944+
"available_categories": list(TOOL_CATEGORIES.keys())
1945+
}
1946+
1947+
return {
1948+
"category": category_filter,
1949+
"tools": TOOL_CATEGORIES[category_filter],
1950+
"tool_count": len(TOOL_CATEGORIES[category_filter])
1951+
}
1952+
else:
1953+
# Return all categories
1954+
result = {
1955+
"categories": {},
1956+
"total_tools": 0
1957+
}
1958+
1959+
for cat_name, tool_list in TOOL_CATEGORIES.items():
1960+
result["categories"][cat_name] = {
1961+
"tools": tool_list,
1962+
"count": len(tool_list)
1963+
}
1964+
result["total_tools"] += len(tool_list)
1965+
1966+
return result
1967+
1968+
17941969
# Register aliases for backward compatibility
1795-
from .tool_registry import ToolRegistration, TOOL_REGISTRY
1970+
from .tool_registry import ToolRegistration
17961971

17971972
# Alias: ARANGO_GRAPH_TRAVERSAL -> handle_traverse
17981973
TOOL_REGISTRY[ARANGO_GRAPH_TRAVERSAL] = ToolRegistration(

‎mcp_arangodb_async/models.py‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -414,3 +414,71 @@ class GraphStatisticsArgs(BaseModel):
414414
class ArangoDatabaseStatusArgs(BaseModel):
415415
"""Arguments for arango_database_status tool (no parameters required)."""
416416
pass
417+
418+
419+
# Pattern 1: Progressive Tool Discovery
420+
class SearchToolsArgs(BaseModel):
421+
"""Arguments for arango_search_tools."""
422+
keywords: List[str] = Field(
423+
description="Keywords to search for in tool names and descriptions"
424+
)
425+
categories: Optional[List[str]] = Field(
426+
default=None,
427+
description="Filter by categories: core_data, indexing, validation, schema, query, graph_basic, graph_advanced, aliases, health"
428+
)
429+
detail_level: Literal["name", "summary", "full"] = Field(
430+
default="summary",
431+
description="Level of detail: 'name' (just names), 'summary' (names + descriptions), 'full' (complete schemas)"
432+
)
433+
434+
435+
class ListToolsByCategoryArgs(BaseModel):
436+
"""Arguments for arango_list_tools_by_category."""
437+
category: Optional[str] = Field(
438+
default=None,
439+
description="Category to filter by. If None, returns all categories with their tools."
440+
)
441+
442+
443+
# Pattern 2: Context Switching
444+
class SwitchContextArgs(BaseModel):
445+
"""Arguments for arango_switch_context."""
446+
context: Literal[
447+
"baseline", "data_analysis", "graph_modeling",
448+
"bulk_operations", "schema_validation", "full"
449+
] = Field(
450+
description="Workflow context to switch to"
451+
)
452+
453+
454+
class GetActiveContextArgs(BaseModel):
455+
"""Arguments for arango_get_active_context."""
456+
pass
457+
458+
459+
class ListContextsArgs(BaseModel):
460+
"""Arguments for arango_list_contexts."""
461+
include_tools: bool = Field(
462+
default=False,
463+
description="Include tool lists for each context"
464+
)
465+
466+
467+
# Pattern 3: Tool Unloading
468+
class AdvanceWorkflowStageArgs(BaseModel):
469+
"""Arguments for arango_advance_workflow_stage."""
470+
stage: Literal["setup", "data_loading", "analysis", "cleanup"] = Field(
471+
description="Workflow stage to advance to"
472+
)
473+
474+
475+
class GetToolUsageStatsArgs(BaseModel):
476+
"""Arguments for arango_get_tool_usage_stats."""
477+
pass
478+
479+
480+
class UnloadToolsArgs(BaseModel):
481+
"""Arguments for arango_unload_tools."""
482+
tool_names: List[str] = Field(
483+
description="List of tool names to unload from active context"
484+
)

‎mcp_arangodb_async/tools.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,3 +104,17 @@
104104

105105
# Database Status Tool
106106
ARANGO_DATABASE_STATUS = "arango_database_status"
107+
108+
# Pattern 1: Progressive Tool Discovery
109+
ARANGO_SEARCH_TOOLS = "arango_search_tools"
110+
ARANGO_LIST_TOOLS_BY_CATEGORY = "arango_list_tools_by_category"
111+
112+
# Pattern 2: Context Switching
113+
ARANGO_SWITCH_CONTEXT = "arango_switch_context"
114+
ARANGO_GET_ACTIVE_CONTEXT = "arango_get_active_context"
115+
ARANGO_LIST_CONTEXTS = "arango_list_contexts"
116+
117+
# Pattern 3: Tool Unloading
118+
ARANGO_ADVANCE_WORKFLOW_STAGE = "arango_advance_workflow_stage"
119+
ARANGO_GET_TOOL_USAGE_STATS = "arango_get_tool_usage_stats"
120+
ARANGO_UNLOAD_TOOLS = "arango_unload_tools"

0 commit comments

Comments
 (0)