Skip to content

adding-tools.md

Latest commit

 

History

History
207 lines (161 loc) · 7.71 KB

adding-tools.md

File metadata and controls

207 lines (161 loc) · 7.71 KB
Raw

Adding Tools

Tools are the MCP primitive this server exposes: Python functions AI assistants discover and call. All tools live in server/tools.py, registered inside load_tools().

The pattern

@mcp_server.tool
def calculate_sum(a: int, b: int) -> dict:
    """
    Calculate the sum of two numbers.

    Args:
        a: First number
        b: Second number

    Returns:
        dict: Contains the sum result
    """
    return {"result": a + b}

Restart the server; the tool is automatically discovered by clients and automatically covered by the integration tests.

Rules that matter

Rule Why
Type hints are mandatory FastMCP derives the validation schema from them
Docstrings are the contract AI assistants read them to decide when to call the tool, so write for the model, not just for humans
Return structured data dicts or Pydantic models, never plain strings
Handle errors gracefully wrap SDK calls in try/except and return {"error": ...}; a raised exception gives the model nothing to work with
Keep tools in server/tools.py one registration point; don't create new tool files
Run uv run pytest tests/ after changes test_call_tools exercises every registered tool

Using the Databricks SDK

Two client factories in server/utils.py; see Architecture → Authentication model for how they resolve per environment:

from server import utils


@mcp_server.tool
def list_clusters() -> dict:
    """List Databricks clusters visible to the app."""
    try:
        w = utils.get_workspace_client()  # app service principal
        return {"clusters": [c.cluster_name for c in w.clusters.list()]}
    except Exception as e:
        return {"success": False, "error": str(e)}


@mcp_server.tool
def whoami() -> dict:
    """Identity of the calling end user."""
    try:
        w = utils.get_user_authenticated_workspace_client()  # end user
        user = w.current_user.me()
        return {"user_name": user.user_name, "active": user.active}
    except Exception as e:
        return {"success": False, "error": str(e)}

The existing get_current_user and ask_genie tools in server/tools.py are real, working references for both auth modes.

Recipe: wrap a Databricks-hosted agent as a tool

This is the accelerator's core move: an agent running in Databricks becomes one MCP tool here, making it discoverable and callable by external agents. ask_genie is the built-in example (it fronts Genie). This recipe covers the general case: an agent you built with Agent Framework / Agent Bricks (or any MLflow ResponsesAgent) deployed on a Model Serving endpoint.

1. Know your endpoint

You need the serving endpoint name and its request format:

  • Responses format (current recommendation, MLflow ResponsesAgent): request body {"input": [{"role": "user", "content": "..."}]}
  • Chat-completions format (older agents): {"messages": [{"role": "user", "content": "..."}]}

Check the endpoint's serving page or query it once from a notebook to confirm.

2. Make the endpoint name configuration, not code

Follow the existing pattern exactly: new env var in app.yaml:

  - name: AGENT_ENDPOINT_NAME
    valueFrom: agent-endpoint

and bind a matching app resource (type Serving endpoint, permission Can query, resource key agent-endpoint) as in Secrets & Configuration, step 4. Locally: export AGENT_ENDPOINT_NAME="<endpoint-name>".

Per repo rule 11, an app.yaml change means updating the deploy skill (.claude/skills/deploy-accelerator/SKILL.md) in the same PR; the resource-binding table in Phase 4 gets a new row.

3. Write the tool

Minimal version using the SDK's raw API client (no new dependencies), in server/tools.py inside load_tools():

import os
from databricks.sdk import WorkspaceClient


@mcp_server.tool
def ask_supply_chain_agent(question: str) -> dict:
    """
    Ask the supply-chain analyst agent a question about inventory,
    suppliers, or logistics. Use this instead of ask_genie when the
    question is about supply-chain operations rather than general data.

    Args:
        question: The user's natural-language question.

    Returns:
        dict with status and the agent's answer, or an error.
    """
    try:
        endpoint = os.getenv("AGENT_ENDPOINT_NAME")
        if not endpoint:
            return {
                "error": "AGENT_ENDPOINT_NAME not set",
                "message": "Set it in app.yaml using valueFrom: agent-endpoint",
            }

        w = WorkspaceClient()  # lazy: created at call time, never at import
        response = w.api_client.do(
            "POST",
            f"/serving-endpoints/{endpoint}/invocations",
            body={"input": [{"role": "user", "content": question}]},
        )

        # ResponsesAgent output: list of items; text lives in
        # output[*].content[*].text; adjust to your agent's schema.
        parts = []
        for item in response.get("output", []):
            for block in item.get("content", []) or []:
                if block.get("type") == "output_text" and block.get("text"):
                    parts.append(block["text"])

        return {
            "status": "ok",
            "question": question,
            "answer": "\n\n".join(parts) if parts else str(response),
        }
    except Exception as e:
        return {"status": "error", "question": question, "error": str(e)}

Alternative: the Databricks OpenAI client (w.serving_endpoints.get_open_ai_client() → client.responses.create(model=endpoint, input=[...])) gives you typed responses and streaming, but it requires adding the openai package to both pyproject.toml and requirements.txt (dependency rule, Phase 0.4 of the deploy skill). Prefer the raw call until you need those features.

Rules that carry the weight here:

  • The docstring is routing logic. External agents choose between your tools by reading it, so say when to use this tool instead of the others (see the example's second sentence).
  • Lazy client, error dicts, per repo rule 12: no auth/network at import time; failures return structured errors.
  • Multi-turn: keep state with the caller, like ask_genie does. If the agent supports threads/conversations, accept an optional id parameter and return it; otherwise accept an optional history the caller replays.

4. Grant, test, deploy

  1. Grant the app's service principal Can query on the serving endpoint (binding the resource in step 2 typically handles this; verify).
  2. uv run pytest tests/: test_call_tools picks the new tool up automatically (it will return the structured "not set" error in the hermetic suite, which passes).
  3. Deploy (/deploy-accelerator) and verify with ./scripts/dev/query_remote.sh, then in AI Playground, confirming the model chooses the right tool for a supply-chain-style question vs. a generic data question. If it picks wrong, fix the docstrings, not the code.

Once deployed, the agent is discoverable and callable by every MCP client (AI Playground, Claude, a Copilot Studio agent in Teams; see the Teams guide) with no client-side changes.

References: Query an agent deployed on Databricks · Query LLMs and agents

Renaming the project

  1. name in pyproject.toml (and the [project.scripts] command name)
  2. FastMCP(name="...") in server/app.py
  3. README references