Optimize MCP tool docstrings for token efficiency by QuentinAmbard · Pull Request #394 · databricks-solutions/ai-dev-kit

QuentinAmbard · 2026-03-30T08:22:16Z

Summary

This PR optimizes MCP tools for better AI efficiency through docstring compression, tool consolidation, and improved documentation.

Key Changes

1. Tool Consolidation (77 → 44 tools, 43% reduction)

Module	Before	After	Consolidated Into
pipelines.py	10	2	`manage_pipeline`, `manage_pipeline_run`
volume_files.py	6	1	`manage_volume_files`
vector_search.py	8	4	`manage_vs_endpoint`, `manage_vs_index`, `query_vs_index`, `manage_vs_data`
genie.py	5	2	`manage_genie`, `ask_genie`
aibi_dashboards.py	4	1	`manage_dashboard`
serving.py	3	1	`manage_serving_endpoint`
apps.py	3	1	`manage_app`
file.py	2	1	`manage_workspace_files`
sql.py	6	5	`manage_warehouse` replaces list/get_best
lakebase.py	8	4	`manage_lakebase_database`, `manage_lakebase_branch`, `manage_lakebase_sync`, `generate_lakebase_credential`

2. Docstring Optimization (~89% reduction in documentation size)

Reduced verbosity across all 18 tool files
Maintained all functional information while being concise
Added skill references to complex tools for detailed documentation

3. Comprehensive Parameter Documentation

Unity Catalog tools: All 9 tools have detailed parameter documentation per action/object_type
Dashboard tool: Added CRITICAL validation steps requiring schema checks and query testing
All consolidated tools: Clear documentation of required params per action

4. Bug Fixes (discovered via integration tests)

Component	Bug	Fix
Apps	`deploy_app` returned `Wait[AppDeployment]` object instead of `AppDeployment`	Use `wait_obj.response` to get actual deployment object
Genie	Used manager methods that don't exist in SDK	Use SDK methods: `w.genie.trash_space()`, `w.genie.update_space()`, `w.genie.create_space()`, `w.genie.get_space()`
Workspace Files	Directory upload failed	Fixed upload for directories in databricks-tools-core
Jobs	Notebook path handling incorrect	Fixed job notebook path handling
Vector Search	Index operations failing	Fixed vector search index operations
Dashboard	Widget handling issues	Fixed dashboard widget handling
Agent Bricks	Manager listing broken	Fixed agent bricks manager listing

5. Integration Test Infrastructure

Added comprehensive integration tests for all MCP tools
Added test runner script with parallel execution and per-suite timeouts
Added fixtures for workspace, catalog, and resource cleanup
Improved test reliability with quota skip handling

Skill References Added

Complex tools now reference corresponding skills for detailed documentation:

databricks-aibi-dashboards - widget JSON structures, filter types
databricks-vector-search - index spec details (delta_sync, direct_access)
agent-bricks - KA/MAS configuration
databricks-genie - Genie space setup
databricks-jobs - task configuration details
databricks-spark-declarative-pipelines - pipeline setup
databricks-lakebase-provisioned/autoscale - database types
databricks-unity-catalog - UC guidance
databricks-model-serving - endpoint configuration
databricks-app-python - app development

Running Integration Tests

# Navigate to the MCP server directory
cd databricks-mcp-server

# Run all integration tests
python tests/integration/run_tests.py --all

# Run specific test suites
python tests/integration/run_tests.py --suites genie,pipelines,jobs

# Run with pytest directly (single suite)
python -m pytest tests/integration/sql -m integration -v --tb=short

# Run non-slow tests only
python -m pytest tests/integration -m "integration and not slow" -v

# Run a specific test
python -m pytest tests/integration/apps/test_apps.py::TestAppLifecycle::test_full_app_lifecycle -v

Prerequisites:

Set DATABRICKS_HOST and DATABRICKS_TOKEN environment variables
Or configure ~/.databrickscfg with a DEFAULT profile

Commits

Optimize MCP tool docstrings - ~89% reduction in documentation size
Add parameter context - Clarify ambiguous docstring params
Consolidate MCP tools - 77 → 44 tools (43% reduction)
Add integration test infrastructure - Comprehensive tests and bug fixes
Fix deploy_app - Handle SDK Wait[AppDeployment] return type
Clarify MCP tool usage - Improve Genie skill documentation
Fix typo - Remove garbage characters from dashboard docstring
Fix genie tools - Use SDK methods instead of manager
Improve test reliability - Add timeouts and quota skip handling
Improve UC tool docstrings - Comprehensive parameter documentation
Add CRITICAL validation - Dashboard tool requires query testing

Test plan

All integration tests pass
Verify AI can still use tools correctly with compressed docstrings
Confirm skill references load correct documentation when needed
Test complex tools (dashboards, vector search, genie) work as expected

🤖 Generated with Claude Code

- Reduce docstring verbosity across all 18 tool files (~89% reduction) - Keep all functional information while being concise - Add skill references to complex tools (dashboards, vector search, genie, jobs, pipelines, lakebase, unity catalog, serving, apps, agent bricks) - Maintain human readability with bullet points and structure - Preserve critical warnings (ASK USER FIRST, CONFIRM WITH USER) - Keep return format hints for AI parsing Net reduction: 1,843 lines across 18 files 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

- agent_bricks.py: Add context for description, instructions, volume_path, examples - genie.py: Add context for table_identifiers, description, sample_questions, serialized_space - jobs.py: Add context for tasks, job_clusters, environments, schedule, git_source - lakebase.py: Add context for source_branch, ttl_seconds, is_protected, autoscaling params, and sync source/target table names - pipelines.py: Add context for root_path, workspace_file_paths, extra_settings, full_refresh 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

Tool consolidations: - pipelines.py: 10→2 (manage_pipeline, manage_pipeline_run) - volume_files.py: 6→1 (manage_volume_files) - aibi_dashboards.py: 4→1 (manage_dashboard) - vector_search.py: 8→4 (manage_vs_endpoint, manage_vs_index, query_vs_index, manage_vs_data) - genie.py: 5→2 (manage_genie, ask_genie) - serving.py: 3→1 (manage_serving_endpoint) - apps.py: 3→1 (manage_app) - file.py: 2→1 (manage_workspace_files) - sql.py: 6→5 (manage_warehouse replaces list/get_best) - lakebase.py: 8→4 (manage_lakebase_database, manage_lakebase_branch, manage_lakebase_sync, generate_lakebase_credential) Key patterns: - All consolidated tools use an action parameter - Each action has required params documented in docstring - Error messages specify which params are required - Hot paths (query_vs_index, ask_genie) kept separate for clarity - All skills updated with action tables and examples 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

Test infrastructure: - Add comprehensive integration tests for all MCP tools - Add test runner script with parallel execution support - Add fixtures for workspace, catalog, and resource cleanup - Add test resources (PDFs, SQL files, app configs) Bug fixes in databricks-tools-core: - Fix workspace file upload for directories - Fix job notebook path handling - Fix vector search index operations - Fix apps API responses - Fix dashboard widget handling - Fix agent bricks manager listing Bug fixes in MCP server tools: - Add quota skip handling for apps test - Fix genie space operations - Fix lakebase database operations - Fix compute cluster lifecycle handling - Fix dashboard operations 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

The Databricks SDK's w.apps.deploy() returns a Wait[AppDeployment] object, not an AppDeployment directly. The previous code passed the Wait object to _deployment_to_dict(), which caused getattr() to return None for all attributes since the Wait object doesn't have them. This fix uses wait_obj.response to get the actual AppDeployment object before converting it to a dictionary. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

QuentinAmbard · 2026-04-01T07:18:06Z

Recommended workflow:

Merge PR Optimize MCP tool docstrings for token efficiency #394 first (MCP tools + apps fix + integration tests)
Rebase PR Enhance AI/BI dashboard skill with comprehensive widget documentation #374 on the updated main - Git will automatically drop the MCP commits since they're already in main
PR Enhance AI/BI dashboard skill with comprehensive widget documentation #374's diff will then only show dashboard skill improvements

- Add tools summary table at top of MCP Tools section - Change code blocks from python syntax to plain text - Add "# MCP Tool: <name>" comments to clarify these are tool calls, not Python code - Move Supporting Tools table to main tools table 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

Remove garbage characters from widget documentation. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

- Use w.genie.trash_space() in _delete_genie_resource - Add _find_space_by_name() using SDK's list_spaces with pagination - Use w.genie.update_space() and w.genie.create_space() for space management - Use w.genie.get_space() with include_serialized_space in _get_genie_space - Fix validation to allow space_id for updates without display_name 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

- Add per-suite timeout in run_tests.py (10 min default, configurable) - Improve apps test with better cleanup and assertions - Add skip logic for quota-exceeded scenarios 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

…cumentation Add detailed parameter documentation to all 9 Unity Catalog MCP tools: - manage_uc_objects: Document parameters by object_type (catalog/schema/volume/function) - manage_uc_grants: Add privilege lists per securable type - manage_uc_storage: Detail credential and external_location parameters - manage_uc_connections: Document connection_type options and create_foreign_catalog - manage_uc_tags: Detail set_tags/unset_tags/query parameters - manage_uc_security_policies: Document row filter and column mask parameters - manage_uc_monitors: Detail monitor creation and refresh parameters - manage_uc_sharing: Document share/recipient/provider resource types - manage_metric_views: Detail dimension/measure format and query parameters 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

Add clear instructions requiring users to: 0. Review the databricks-aibi-dashboards skill for widget JSON structure 1. Call get_table_stats_and_schema() for table schemas 2. Call execute_sql() to test EVERY query before use This prevents widgets from showing errors due to untested queries. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

Quentin Ambard and others added 5 commits March 30, 2026 10:20

QuentinAmbard mentioned this pull request Apr 1, 2026

Enhance AI/BI dashboard skill with comprehensive widget documentation #374

Open

3 tasks

Quentin Ambard and others added 6 commits April 1, 2026 09:23

Fix typo in aibi_dashboards.py docstring

fe7ea10

Remove garbage characters from widget documentation. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Optimize MCP tool docstrings for token efficiency#394

Optimize MCP tool docstrings for token efficiency#394
QuentinAmbard wants to merge 11 commits intomainfrom
mcp-tool-token-optimization

QuentinAmbard commented Mar 30, 2026 •

edited

Loading

Uh oh!

QuentinAmbard commented Apr 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

QuentinAmbard commented Mar 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Key Changes

Skill References Added

Running Integration Tests

Commits

Test plan

Uh oh!

QuentinAmbard commented Apr 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

QuentinAmbard commented Mar 30, 2026 •

edited

Loading