databricks-agentbricks 0.2.0__py3-none-any.whl

databricks_agentkit/client.py

@@ -0,0 +1,31 @@
+"""Ergonomic resource wrapper over the generated service API."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+from databricks_agentkit._api_client import _AgentBricksApiClient
+from databricks_agentkit.memory_store import MemoryStores
+from databricks_agentkit.session_store import SessionStores
+
+if TYPE_CHECKING:
+    # Imported lazily elsewhere: pulling databricks.sdk costs ~0.7s, so keeping it out of the
+    # import path lets local CLI commands (help, init, tools, bind, completion) start fast.
+    from databricks.sdk import WorkspaceClient
+
+
+class AgentKitClient:
+    """High-level client for AgentKit memory and session APIs.
+
+    ``WorkspaceClient.mason`` will replace the private transport once the generated SDK is
+    released; the resource-oriented public surface remains unchanged.
+
+    Args:
+        workspace_client: An authenticated Databricks workspace client. When omitted,
+            the Databricks SDK's default authentication resolution is used.
+    """
+
+    def __init__(self, workspace_client: Optional[WorkspaceClient] = None) -> None:
+        api = _AgentBricksApiClient(workspace_client=workspace_client)
+        self.memory_stores = MemoryStores(api)
+        self.session_stores = SessionStores(api)

databricks_agentkit/langgraph/__init__.py

@@ -0,0 +1,98 @@
+"""LangGraph adapter for running an agent on Databricks (installed via ``databricks-agentbricks[langgraph]``).
+
+Composable pieces you drop into an existing LangGraph agent — a session-store checkpointer, MCP
+tools declared in ``agent.toml``, long-term memory tools, and MLflow tracing. Each maps onto a slot
+LangGraph already has, so migrating an existing agent is a graft, not a rewrite::
+
+    from databricks_agentkit.langgraph import (
+        checkpointer,
+        thread_config,
+        mcp_tools,
+        configure_tracing,
+    )
+
+    configure_tracing()
+    agent = create_agent(
+        model=...,
+        tools=[
+            *your_tools,
+            *await mcp_tools(),
+        ],  # mcp_tools() includes your agent.toml servers
+        checkpointer=checkpointer(),  # durable when AGENT_SESSION_STORE is set
+    )
+    result = await agent.ainvoke(inputs, config=thread_config(session_id))
+
+These need the agent stack (databricks-langchain, langgraph, langchain, fastapi, mlflow), so they sit
+behind the ``[langgraph]`` extra to keep a plain ``databricks-agentbricks`` installation independent of agent frameworks.
+
+``__all__`` is the curated surface. Other entry points (``mcp_client``, ``DatabricksSessionStoreSaver``)
+are reachable by their submodule paths but not re-exported here.
+"""
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from databricks_agentkit.langgraph.genie import genie_tools
+    from databricks_agentkit.langgraph.mcp import mcp_tools
+    from databricks_agentkit.langgraph.memory import memory_tools
+    from databricks_agentkit.langgraph.session_store import checkpointer, thread_config
+    from databricks_agentkit.runtime import (
+        start_trace,
+        workspace_client,
+        workspace_headers,
+    )
+
+
+def configure_tracing() -> None:
+    """Enable MLflow tracing with LangChain autologging. Call once at startup.
+
+    Safe to call unconditionally — tracing turns on only when the MLflow destination and experiment
+    are configured in the environment (see :func:`databricks_agentkit.runtime.configure_tracing`).
+    """
+    import mlflow
+
+    from databricks_agentkit.runtime import configure_tracing as _configure_tracing
+
+    _configure_tracing(autolog=mlflow.langchain.autolog)
+
+
+__all__ = [
+    "genie_tools",
+    # MCP tools from agent.toml (plus any servers you pass) — add them to your agent's tool list.
+    "mcp_tools",
+    # Long-term memory tools (opt-in via AGENT_MEMORY_STORE) — add to your tool list.
+    "memory_tools",
+    # Session persistence — pass checkpointer() to create_agent(checkpointer=...) and
+    # thread_config(session_id) as the per-request run config.
+    "checkpointer",
+    "thread_config",
+    # MLflow tracing (LangChain autolog bound in) — call configure_tracing() once at startup.
+    "configure_tracing",
+    # Wrap each invocation in start_trace() so a trace is recorded (pass session_id= to tag it).
+    "start_trace",
+    # Workspace SDK client construction.
+    "workspace_client",
+    "workspace_headers",
+]
+
+# Re-exports resolved lazily (PEP 562) so importing one submodule (e.g. ``.mcp``) does not eagerly
+# pull in the others' dependencies. ``configure_tracing`` is defined above (binds LangChain autolog).
+_MODULE_BY_NAME = {
+    "genie_tools": "databricks_agentkit.langgraph.genie",
+    "mcp_tools": "databricks_agentkit.langgraph.mcp",
+    "memory_tools": "databricks_agentkit.langgraph.memory",
+    "checkpointer": "databricks_agentkit.langgraph.session_store",
+    "thread_config": "databricks_agentkit.langgraph.session_store",
+    "start_trace": "databricks_agentkit.runtime",
+    "workspace_client": "databricks_agentkit.runtime",
+    "workspace_headers": "databricks_agentkit.runtime",
+}
+
+
+def __getattr__(name: str) -> object:
+    module = _MODULE_BY_NAME.get(name)
+    if module is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    import importlib
+
+    return getattr(importlib.import_module(module), name)

databricks_agentkit/langgraph/genie.py

@@ -0,0 +1,50 @@
+"""Native Genie Agent tools for LangGraph, configured through ``agent.toml``."""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+from langchain_core.tools import BaseTool, tool
+
+from databricks_agentkit.runtime.auth import AuthError
+from databricks_agentkit.runtime.genie import GenieAgent
+from databricks_agentkit.runtime.tool_manifest import load_tools
+
+if TYPE_CHECKING:
+    from databricks.sdk import WorkspaceClient
+
+
+def genie_tools(
+    *, workspace_client_for: Callable[[str], WorkspaceClient] | None = None
+) -> list[BaseTool]:
+    """Build ask, poll and query-result tools for each configured Genie Agent.
+
+    Space IDs are bound from the manifest, not exposed as model arguments. Construction performs
+    no authentication or network requests; an agent with no Genie Agent bindings gets no tools.
+    """
+    tools = []
+    for record in load_tools(expected_framework="langgraph"):
+        if record.kind != "genie_agent":
+            continue
+        if (
+            record.auth == "user"
+            and workspace_client_for is None
+            and os.getenv("DATABRICKS_APP_NAME")
+        ):
+            raise AuthError(
+                "MCP_USER_AUTHORIZATION_MISSING",
+                "This deployed Genie integration requires request-user authorization.",
+                integration_id=record.id,
+            )
+        agent = GenieAgent(
+            record.space_id or "",
+            auth=record.auth or "app",
+            workspace_client_for=workspace_client_for,
+        )
+        tools.extend(
+            tool(f"{record.id}_{method.__name__}")(method)
+            for method in (agent.ask, agent.poll, agent.query_result)
+        )
+    return tools

databricks_agentkit/langgraph/mcp.py

@@ -0,0 +1,216 @@
+"""Build MCP tools for the agent from the servers declared in ``agent.toml`` (plus any the agent adds).
+
+``mcp_tools()`` is the entry point: it reads the MCP servers declared in ``agent.toml``
+(sandbox/mcp + uc_function), fetches their LangChain tools with sandbox downscoping applied, and
+returns them. An agent with its own hand-built servers passes them as ``extra_servers``; leaving
+``agent.toml`` empty simply yields no declared servers. Typical agent use::
+
+    tools = await mcp_tools()  # just the agent.toml servers
+    tools = await mcp_tools(build_mcp_servers())  # agent.toml servers + the agent's own
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import TYPE_CHECKING, Any
+
+from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient
+from langchain_mcp_adapters.sessions import create_session
+
+if TYPE_CHECKING:
+    from databricks_langchain import MCPServer
+
+from databricks_agentkit.runtime import mcp_auth
+from databricks_agentkit.runtime.auth import AuthError
+from databricks_agentkit.runtime.tool_manifest import ToolRecord, downscope_wire, load_tools
+from databricks_agentkit.runtime.workspace import workspace_client, workspace_headers
+
+logger = logging.getLogger(__name__)
+_auth_error = mcp_auth.mcp_auth_error
+_tool_error = mcp_auth.mcp_tool_error
+
+
+def _server_from_tool(
+    tool: ToolRecord,
+    *,
+    workspace_client_for: mcp_auth.WorkspaceClientResolver | None = None,
+) -> DatabricksMCPServer | None:
+    if tool.kind not in {"sandbox", "mcp", "uc_function", "genie_one"}:
+        return None
+    mode = tool.auth or "app"
+    client = mcp_auth.resolve_mcp_workspace_client(
+        mode, tool.id, workspace_client_for, workspace_client
+    )
+    host = client.config.host.rstrip("/")
+    if tool.kind in {"sandbox", "mcp", "genie_one"}:
+        url = (
+            f"{host}/api/2.0/mcp/genie"
+            if tool.kind == "genie_one"
+            else f"{host}/ai-gateway/mcp-services/{tool.service}"
+        )
+        return DatabricksMCPServer(
+            name=tool.id,
+            url=url,
+            headers=workspace_headers() or None,
+            workspace_client=client,
+            timeout=120.0,
+        )
+    if tool.kind == "uc_function":
+        catalog, schema, function_name = (tool.function or "").split(".")
+        return DatabricksMCPServer.from_uc_function(
+            catalog=catalog,
+            schema=schema,
+            function_name=function_name,
+            name=tool.id,
+            headers=workspace_headers() or None,
+            workspace_client=client,
+            timeout=120.0,
+        )
+    return None
+
+
+def _declared_servers(
+    *,
+    workspace_client_for: mcp_auth.WorkspaceClientResolver | None = None,
+) -> list[DatabricksMCPServer]:
+    """The MCP servers declared in the agent's ``agent.toml`` (may be empty)."""
+    tools = load_tools(expected_framework="langgraph")
+    return [
+        server
+        for tool in tools
+        if (server := _server_from_tool(tool, workspace_client_for=workspace_client_for))
+        is not None
+    ]
+
+
+def _sandbox_interceptor(
+    tools: tuple[ToolRecord, ...],
+    *,
+    workspace_client_for: mcp_auth.WorkspaceClientResolver | None = None,
+):
+    declared = {tool.id: tool for tool in tools}
+
+    async def interceptor(request: Any, handler: Any) -> Any:
+        tool = declared.get(request.server_name)
+        if tool is None:
+            return await handler(request)
+        request_user = tool.auth == "user"
+        try:
+            if tool.kind == "sandbox":
+                server = _server_from_tool(tool, workspace_client_for=workspace_client_for)
+                if server is None:
+                    raise RuntimeError(f"Could not build sandbox MCP server {tool.id!r}.")
+                async with create_session(server.to_connection_dict()) as session:
+                    await session.initialize()
+                    result = await session.call_tool(
+                        request.name,
+                        request.args,
+                        meta={"downscope": downscope_wire(tool)},
+                    )
+            else:
+                result = await handler(request)
+        except Exception as error:
+            if request_user:
+                raise _auth_error(error, tool.id) or AuthError(
+                    "MCP_TOOL_FAILED", "The configured MCP tool failed.", 502, tool.id
+                ) from None
+            raise
+        if (
+            request_user
+            and getattr(result, "isError", False)
+            and (error := _tool_error(result, tool.id))
+        ):
+            raise error
+        return result
+
+    return interceptor
+
+
+def mcp_client(
+    servers: list[DatabricksMCPServer],
+    *,
+    workspace_client_for: mcp_auth.WorkspaceClientResolver | None = None,
+    tools: tuple[ToolRecord, ...] | None = None,
+) -> DatabricksMultiServerMCPClient:
+    """A multi-server MCP client over ``servers`` with the sandbox downscoping interceptor attached.
+
+    The interceptor is derived from the ``agent.toml`` manifest, so sandbox tools run downscoped
+    regardless of how the caller drives the returned client (``get_tools`` or otherwise). Callers who
+    build their own client instead take on applying downscoping themselves.
+    """
+    snapshot = tuple(load_tools(expected_framework="langgraph")) if tools is None else tools
+    interceptors = (
+        [_sandbox_interceptor(snapshot, workspace_client_for=workspace_client_for)]
+        if snapshot
+        else []
+    )
+    # DatabricksMCPServer is a subclass of MCPServer, so coerce the type for the API
+    servers_as_mcp: list[MCPServer] = servers  # type: ignore[name-defined,assignment]
+    return DatabricksMultiServerMCPClient(
+        servers_as_mcp, tool_interceptors=interceptors, handle_tool_errors=True
+    )
+
+
+async def mcp_tools(
+    extra_servers: list[DatabricksMCPServer] | None = None,
+    *,
+    workspace_client_for: mcp_auth.WorkspaceClientResolver | None = None,
+) -> list:
+    """Fetch declared tools with request identity and protected sandbox downscoping.
+
+    Includes the MCP servers declared in ``agent.toml``; pass ``extra_servers`` to add servers the
+    agent builds itself. Request-user failures propagate; App, legacy, and optional customer
+    servers retain their existing identity and per-server best-effort discovery behavior.
+    """
+    snapshot = tuple(load_tools(expected_framework="langgraph"))
+    servers = [
+        server
+        for tool in snapshot
+        if (server := _server_from_tool(tool, workspace_client_for=workspace_client_for))
+        is not None
+    ]
+    result = []
+    if servers:
+        client = mcp_client(servers, workspace_client_for=workspace_client_for, tools=snapshot)
+        request_user = {tool.id for tool in snapshot if tool.auth == "user"}
+
+        async def fetch_declared(server: DatabricksMCPServer) -> list:
+            try:
+                return await client.get_tools(server_name=server.name)
+            except Exception as error:
+                if server.name in request_user:
+                    raise _auth_error(error, server.name) or AuthError(
+                        "MCP_TOOL_FAILED",
+                        "Could not discover configured MCP tools.",
+                        502,
+                        server.name,
+                    ) from None
+                logger.warning(
+                    "Failed to fetch MCP tools from server %r; continuing without it.", server.name
+                )
+                return []
+
+        groups = await asyncio.gather(
+            *(fetch_declared(server) for server in servers), return_exceptions=True
+        )
+        for group in groups:
+            if isinstance(group, BaseException):
+                raise group
+            result.extend(group)
+    if extra_servers:
+        optional_client = mcp_client(extra_servers, tools=())
+
+        async def fetch_optional(server: DatabricksMCPServer) -> list:
+            try:
+                return await optional_client.get_tools(server_name=server.name)
+            except Exception:
+                logger.warning(
+                    "Failed to fetch optional MCP tools from server %r; continuing without it.",
+                    server.name,
+                )
+                return []
+
+        for group in await asyncio.gather(*(fetch_optional(server) for server in extra_servers)):
+            result.extend(group)
+    return result

databricks_agentkit/langgraph/memory.py

@@ -0,0 +1,67 @@
+"""Long-term memory tools — opt-in, gated on ``AGENT_MEMORY_STORE``.
+
+Unlike the session store (short-term transcript for one conversation), long-term memory is exposed
+to the model as two tools — ``remember`` and ``recall`` — over the Databricks managed memory store's
+``2.0/agents`` entries API. Facts persist across conversations.
+
+``memory_tools(actor)`` returns the tools when ``AGENT_MEMORY_STORE`` is set, else an empty list, so
+the model never sees them when memory is unconfigured. ``create_agent_graph`` composes them into its
+tool list. This is a stand-in for a future ``databricks-langchain`` memory helper — when the SDK
+provides one, swap the import in ``agent.py`` and delete this file.
+
+Memory entries are per-actor. ``actor`` is the identity whose memory these tools read and write — the
+caller (the agent server) supplies it, typically the signed-in user, so each user gets their own
+long-term memory. It is **closed over**, not a tool argument: the model only ever sees ``remember``'s
+and ``recall``'s own parameters and cannot set or spoof the actor. Pass a fixed value for one shared
+memory (e.g. a team knowledge base).
+"""
+
+from langchain_core.tools import BaseTool, tool
+
+from databricks_agentkit.runtime.workspace import workspace_client
+
+_AGENTS_API = "/api/2.0/agents"
+
+
+def _api():
+    # Build the client lazily (needs workspace auth) so importing this module stays cheap.
+    return workspace_client().api_client
+
+
+def memory_tools(actor: str, store: str | None = None) -> list[BaseTool]:
+    """The long-term-memory tools for ``actor`` when a memory store is configured, else none.
+
+    The store resolves ``store`` arg → ``AGENT_MEMORY_STORE`` env → the ``[memory_store]`` binding in
+    agent.toml (`ab memory bind`) → none (no tools). ``actor`` partitions the store; it is captured
+    in the tools' closures (not exposed to the model). Call this per request with the identity whose
+    memory to use (e.g. the signed-in user).
+    """
+    from databricks_agentkit.runtime.tool_manifest import resolve_memory_store
+
+    store = resolve_memory_store(store)
+    if not store:
+        return []
+    store_path = f"{_AGENTS_API}/memory-stores/{store}"
+
+    @tool
+    def remember(fact: str, topic: str) -> str:
+        """Persist a durable fact about the user in long-term memory."""
+        _api().do(
+            "POST",
+            f"{store_path}/entries",
+            body={"actor_id": actor, "path": f"/{topic}/{fact[:8]}.md", "content": fact},
+        )
+        return "stored"
+
+    @tool
+    def recall(query: str) -> str:
+        """Search the user's long-term memory for facts relevant to the query."""
+        data = _api().do(
+            "POST",
+            f"{store_path}/entries:search",
+            body={"actor_id": actor, "query": query, "limit": 5},
+        )
+        entries = data.get("managed_memory_entries") or []
+        return "\n".join(f"- {e.get('content')}" for e in entries) or "No relevant memories."
+
+    return [remember, recall]