databricks-agentbricks 0.2.0__py3-none-any.whl

databricks_agentkit/_api_client.py

@@ -0,0 +1,673 @@
+"""Private transport for the managed agent store APIs.
+
+The public SDK is the resource-oriented :class:`databricks_agentkit.AgentKitClient`.
+This module temporarily owns the one-method-per-endpoint transport used by that
+wrapper and the CLI. It can be replaced by the generated ``WorkspaceClient.mason``
+service without changing the public resource surface.
+"""
+
+from __future__ import annotations
+
+import configparser
+import os
+import pathlib
+import time
+from typing import TYPE_CHECKING, Any, Optional
+from urllib.parse import quote
+
+from databricks_agentbricks.errors import TRANSIENT_ERROR_CODES, AgentCliError, wrap_api_error
+from databricks_agentkit import models
+
+if TYPE_CHECKING:
+    from databricks.sdk import WorkspaceClient
+
+_BASE = "/api/2.0/agents"
+_MCP_SERVICES_PATH = "/api/2.1/unity-catalog/mcp-services"
+
+# Transient backend failures (e.g. a CANCELLED RPC) usually clear on a retry, so retry safe
+# requests once before surfacing them. Mutating requests must opt in explicitly: their first
+# attempt may have committed even when its response was lost.
+_MAX_ATTEMPTS = 2
+_RETRY_BASE_DELAY_S = 0.2
+
+# Bound the SDK's own retry budget (default 300s). The SDK retries every HTTP 429/503 as platform
+# throttling, so an interactive command otherwise hangs up to 5 minutes on a sustained 429 before
+# surfacing anything; cap it so throttling fails fast with a clear error.
+_CLI_RETRY_TIMEOUT_S = 60
+
+
+def _query(**kwargs: Any) -> dict[str, Any]:
+    """Build a query dict, dropping None and empty values."""
+    return {k: v for k, v in kwargs.items() if v is not None and v != ""}
+
+
+def _body(**kwargs: Any) -> dict[str, Any]:
+    """Build a request body while retaining meaningful empty values."""
+    return {k: v for k, v in kwargs.items() if v is not None}
+
+
+def _as(cls: type, resp: Any) -> Any:
+    """Wrap a JSON response in a typed model, passing non-dicts through unchanged."""
+    return cls(resp) if isinstance(resp, dict) else resp
+
+
+def memory_store_path(name: str) -> str:
+    """Normalize a store id or name into the `memory-stores/{id}` resource segment.
+
+    Validate locally so malformed resource names do not produce misleading endpoint errors.
+    """
+    raw = (name or "").strip()
+    if raw.startswith("memory-stores/"):
+        raw = raw[len("memory-stores/") :]
+    raw = raw.strip().strip("/")
+    if not raw:
+        raise AgentCliError("A memory store id or resource name is required.")
+    if "/" in raw:
+        raise AgentCliError(f"Invalid memory store id or resource name: {name!r}")
+    return f"memory-stores/{raw}"
+
+
+def session_store_path(name: str) -> str:
+    """Normalize a session store name into the `session-stores/{name}` resource segment."""
+    raw = (name or "").strip()
+    if raw.startswith("session-stores/"):
+        raw = raw[len("session-stores/") :]
+    raw = raw.strip().strip("/")
+    if not raw:
+        raise AgentCliError("A session store name is required.")
+    return f"session-stores/{raw}"
+
+
+def memory_entry_path(store: str, entry: str) -> str:
+    entry = (entry or "").strip().strip("/")
+    if entry.startswith("memory-stores/"):
+        return entry
+    if not entry:
+        raise AgentCliError("A memory entry id or resource name is required.")
+    return f"{memory_store_path(store)}/entries/{entry}"
+
+
+def _profile_host(profile: str) -> Optional[str]:
+    config_path = pathlib.Path(
+        os.getenv("DATABRICKS_CONFIG_FILE", pathlib.Path.home() / ".databrickscfg")
+    )
+    parser = configparser.ConfigParser()
+    try:
+        parser.read(config_path)
+    except (OSError, configparser.Error):
+        return None
+    return parser.get(profile, "host", fallback=None)
+
+
+def _bound_retry_timeout(client: WorkspaceClient) -> WorkspaceClient:
+    # Shorten the SDK's default 300s retry budget so 429/503 throttling fails in ~1 min, not ~5.
+    try:
+        client.api_client._api_client._retry_timeout_seconds = _CLI_RETRY_TIMEOUT_S
+    except AttributeError:
+        pass
+    return client
+
+
+def _workspace_client(profile: Optional[str]) -> WorkspaceClient:
+    # Imported here (not at module top) so the ~0.7s databricks.sdk import is paid only when a
+    # command actually builds a client, not on every CLI invocation.
+    from databricks.sdk import WorkspaceClient
+
+    client = _bound_retry_timeout(WorkspaceClient(profile=profile))
+    if not profile or not client.config.workspace_id:
+        return client
+
+    configured_host = _profile_host(profile)
+    resolved_host = client.config.host
+    if not configured_host or configured_host.rstrip("/") == (resolved_host or "").rstrip("/"):
+        return client
+
+    workspace_id = str(client.config.workspace_id)
+    return _bound_retry_timeout(
+        WorkspaceClient(
+            profile=profile,
+            host=configured_host,
+            custom_headers={"X-Databricks-Org-Id": workspace_id},
+        )
+    )
+
+
+class _AgentBricksApiClient:
+    """Private transport for the 2.0 agents API until the generated SDK is available."""
+
+    def __init__(
+        self,
+        profile: Optional[str] = None,
+        *,
+        workspace_client: Optional[WorkspaceClient] = None,
+    ) -> None:
+        if profile is not None and workspace_client is not None:
+            raise ValueError("profile and workspace_client are mutually exclusive")
+        try:
+            self._w = workspace_client or _workspace_client(profile)
+        except Exception as exc:  # noqa: BLE001 - surfaced as a clean CLI error
+            raise AgentCliError(
+                f"Could not initialize Databricks auth: {exc}",
+                hint="Select an existing profile with `ab --profile <name> <command>` "
+                "or authenticate and save it with `ab login --profile <name>`.",
+            ) from exc
+
+    @property
+    def host(self) -> str:
+        return self._w.config.host or "unknown"
+
+    @property
+    def current_user(self) -> str:
+        """The authenticated user's name (used to derive the app source workspace path)."""
+        return str(self._w.current_user.me().user_name or "unknown")
+
+    def ensure_workspace_dir(self, path: str) -> None:
+        """Create a workspace directory (and parents), idempotently.
+
+        MLflow's create_experiment won't create the intermediate folder for a nested experiment
+        path, so callers make the parent dir first.
+        """
+        self._w.workspace.mkdirs(path)
+
+    def create_runtime_store(
+        self,
+        runtime_store_id: str,
+        app_service_principal_id: str,
+        *,
+        app_name: str,
+        retry_transient: bool = False,
+    ) -> models.RuntimeStore:
+        """Create the deployment's Runtime Store through Conversation Store."""
+        return _as(
+            models.RuntimeStore,
+            self._do(
+                "POST",
+                f"{_BASE}/runtime-stores",
+                query={"runtime_store_id": runtime_store_id},
+                body={
+                    "owner": {
+                        "app": {
+                            "name": app_name,
+                            "service_principal_id": app_service_principal_id,
+                        }
+                    }
+                },
+                safe_to_retry=retry_transient,
+            ),
+        )
+
+    def get_runtime_store(self, runtime_store_id: str) -> models.RuntimeStore:
+        """Resolve the service-managed backend and app owner before reusing a store."""
+        return _as(
+            models.RuntimeStore,
+            self._do("GET", f"{_BASE}/runtime-stores/{quote(runtime_store_id, safe='')}"),
+        )
+
+    def delete_runtime_store(self, runtime_store_id: str) -> dict:
+        """Delete a deployment's Runtime Store and its dedicated database."""
+        return self._do(
+            "DELETE",
+            f"{_BASE}/runtime-stores/{quote(runtime_store_id, safe='')}",
+            safe_to_retry=True,
+        )
+
+    def _do(
+        self,
+        method: str,
+        path: str,
+        *,
+        query: Optional[dict] = None,
+        body: Optional[dict] = None,
+        safe_to_retry: bool = False,
+    ) -> Any:
+        retry_allowed = method == "GET" or safe_to_retry
+        delay = _RETRY_BASE_DELAY_S
+        for attempt in range(1, _MAX_ATTEMPTS + 1):
+            try:
+                return self._w.api_client.do(method, path, query=query, body=body)
+            except Exception as exc:  # noqa: BLE001 - normalized to AgentCliError
+                retryable = (
+                    retry_allowed and getattr(exc, "error_code", None) in TRANSIENT_ERROR_CODES
+                )
+                if not retryable or attempt == _MAX_ATTEMPTS:
+                    raise wrap_api_error(exc) from exc
+                time.sleep(delay)
+                delay *= 2
+
+    # --- Unity Catalog MCP Services -----------------------------------------
+
+    def get_mcp_service(self, service: str) -> dict:
+        """Look up a managed MCP service visible to the authenticated user."""
+        return self._do("GET", f"{_MCP_SERVICES_PATH}/{quote(service, safe='')}")
+
+    def list_mcp_services(
+        self, schema: str = "system.ai", page_token: Optional[str] = None
+    ) -> dict:
+        """List MCP Services visible to the user in a Unity Catalog schema."""
+        return self._do(
+            "GET",
+            _MCP_SERVICES_PATH,
+            query=_query(parent=f"schemas/{schema}", page_token=page_token),
+        )
+
+    # --- memory stores -------------------------------------------------------
+
+    def create_memory_store(
+        self,
+        display_name: str,
+        description: Optional[str] = None,
+        *,
+        retry_transient: bool = False,
+    ) -> models.MemoryStore:
+        body = _body(display_name=display_name, description=description)
+        return _as(
+            models.MemoryStore,
+            self._do(
+                "POST",
+                f"{_BASE}/memory-stores",
+                query={"managed_memory_store_id": display_name},
+                body=body,
+                safe_to_retry=retry_transient,
+            ),
+        )
+
+    def get_memory_store(self, name: str) -> models.MemoryStore:
+        return _as(models.MemoryStore, self._do("GET", f"{_BASE}/{memory_store_path(name)}"))
+
+    def list_memory_stores(
+        self, page_size: Optional[int] = None, page_token: Optional[str] = None
+    ) -> models.MemoryStoreList:
+        return _as(
+            models.MemoryStoreList,
+            self._do(
+                "GET",
+                f"{_BASE}/memory-stores",
+                query=_query(page_size=page_size, page_token=page_token),
+            ),
+        )
+
+    def update_memory_store(
+        self, name: str, display_name: Optional[str] = None, description: Optional[str] = None
+    ) -> models.MemoryStore:
+        body = _body(display_name=display_name, description=description)
+        if not body:
+            raise AgentCliError("No fields to update. Provide a display name and/or description.")
+        mask = ",".join(body.keys())
+        return _as(
+            models.MemoryStore,
+            self._do(
+                "PATCH",
+                f"{_BASE}/{memory_store_path(name)}",
+                query=_query(update_mask=mask),
+                body=body,
+            ),
+        )
+
+    def delete_memory_store(self, name: str) -> dict:
+        return self._do("DELETE", f"{_BASE}/{memory_store_path(name)}")
+
+    # --- memory entries ------------------------------------------------------
+
+    def create_memory_entry(
+        self,
+        store: str,
+        actor_id: str,
+        path: str,
+        content: Optional[str] = None,
+        description: Optional[str] = None,
+        session_id: Optional[str] = None,
+        source_type: Optional[str] = None,
+        write_mode: Optional[str] = None,
+    ) -> models.MemoryEntry:
+        body = _body(
+            actor_id=actor_id,
+            path=path,
+            content=content,
+            description=description,
+            session_id=session_id,
+            source_type=source_type,
+            write_mode=write_mode,
+        )
+        return _as(
+            models.MemoryEntry,
+            self._do("POST", f"{_BASE}/{memory_store_path(store)}/entries", body=body),
+        )
+
+    def get_memory_entry(
+        self, store: str, entry: str, read_mask: Optional[str] = None
+    ) -> models.MemoryEntry:
+        return _as(
+            models.MemoryEntry,
+            self._do(
+                "GET",
+                f"{_BASE}/{memory_entry_path(store, entry)}",
+                query=_query(read_mask=read_mask),
+            ),
+        )
+
+    def list_memory_entries(
+        self,
+        store: str,
+        actor_id: str,
+        path_prefix: Optional[str] = None,
+        session_id: Optional[str] = None,
+        page_size: Optional[int] = None,
+        page_token: Optional[str] = None,
+        read_mask: Optional[str] = None,
+    ) -> models.MemoryEntryList:
+        return _as(
+            models.MemoryEntryList,
+            self._do(
+                "GET",
+                f"{_BASE}/{memory_store_path(store)}/entries",
+                query=_query(
+                    actor_id=actor_id,
+                    path_prefix=path_prefix,
+                    session_id=session_id,
+                    page_size=page_size,
+                    page_token=page_token,
+                    read_mask=read_mask,
+                ),
+            ),
+        )
+
+    def search_memory_entries(
+        self,
+        store: str,
+        actor_id: str,
+        query: str,
+        limit: Optional[int] = None,
+        page_size: Optional[int] = None,
+        path_prefix: Optional[str] = None,
+        session_id: Optional[str] = None,
+        read_mask: Optional[str] = None,
+    ) -> models.MemorySearchResult:
+        body = _body(
+            actor_id=actor_id,
+            query=query,
+            limit=limit,
+            page_size=page_size,
+            path_prefix=path_prefix,
+            session_id=session_id,
+            read_mask=read_mask,
+        )
+        return _as(
+            models.MemorySearchResult,
+            self._do(
+                "POST",
+                f"{_BASE}/{memory_store_path(store)}/entries:search",
+                body=body,
+                safe_to_retry=True,
+            ),
+        )
+
+    def update_memory_entry(
+        self,
+        store: str,
+        entry: str,
+        content: Optional[str] = None,
+        description: Optional[str] = None,
+    ) -> models.MemoryEntry:
+        body = _body(content=content, description=description)
+        if not body:
+            raise AgentCliError("No fields to update. Provide content and/or a description.")
+        mask = ",".join(body.keys())
+        return _as(
+            models.MemoryEntry,
+            self._do(
+                "PATCH",
+                f"{_BASE}/{memory_entry_path(store, entry)}",
+                query=_query(update_mask=mask),
+                body=body,
+            ),
+        )
+
+    def delete_memory_entry(self, store: str, entry: str) -> dict:
+        return self._do("DELETE", f"{_BASE}/{memory_entry_path(store, entry)}")
+
+    # --- session stores ------------------------------------------------------
+
+    def create_session_store(
+        self,
+        name: str,
+        description: Optional[str] = None,
+        metadata: Optional[dict] = None,
+        *,
+        retry_transient: bool = False,
+    ) -> models.SessionStore:
+        body = _body(description=description, metadata=metadata)
+        return _as(
+            models.SessionStore,
+            self._do(
+                "POST",
+                f"{_BASE}/session-stores",
+                query={"session_store_name": name, "session_store_id": name},
+                body=body,
+                safe_to_retry=retry_transient,
+            ),
+        )
+
+    def get_session_store(self, name: str) -> models.SessionStore:
+        return _as(models.SessionStore, self._do("GET", f"{_BASE}/{session_store_path(name)}"))
+
+    def list_session_stores(
+        self, page_size: Optional[int] = None, page_token: Optional[str] = None
+    ) -> models.SessionStoreList:
+        return _as(
+            models.SessionStoreList,
+            self._do(
+                "GET",
+                f"{_BASE}/session-stores",
+                query=_query(page_size=page_size, page_token=page_token),
+            ),
+        )
+
+    def update_session_store(
+        self, name: str, description: Optional[str] = None, metadata: Optional[dict] = None
+    ) -> models.SessionStore:
+        body = _body(description=description, metadata=metadata)
+        if not body:
+            raise AgentCliError("No fields to update. Provide a description and/or metadata.")
+        mask = ",".join(body.keys())
+        return _as(
+            models.SessionStore,
+            self._do(
+                "PATCH",
+                f"{_BASE}/{session_store_path(name)}",
+                query=_query(update_mask=mask),
+                body=body,
+            ),
+        )
+
+    def delete_session_store(self, name: str) -> dict:
+        return self._do("DELETE", f"{_BASE}/{session_store_path(name)}")
+
+    # --- store permission grants --------------------------------------------
+
+    def grant_session_store_permission(
+        self, name: str, principal_client_id: str, permission: str = "WRITE"
+    ) -> dict:
+        """Grant a service principal READ/WRITE on a session store.
+
+        The store service performs the underlying Lakebase role provisioning and GRANTs itself, so
+        this succeeds without the caller owning the store or holding MANAGE on its Lakebase.
+        """
+        return self._do(
+            "POST",
+            f"{_BASE}/{session_store_path(name)}/permissions:grant",
+            body={
+                "principal": {"type": "SERVICE_PRINCIPAL", "name": principal_client_id},
+                "permission": permission,
+            },
+            safe_to_retry=True,
+        )
+
+    def grant_memory_store_permission(
+        self, name: str, principal_client_id: str, permission: str = "WRITE"
+    ) -> dict:
+        """Grant a service principal READ/WRITE on a memory store (``name`` is its resource id)."""
+        return self._do(
+            "POST",
+            f"{_BASE}/{memory_store_path(name)}/permissions:grant",
+            body={
+                "principal": {"type": "SERVICE_PRINCIPAL", "name": principal_client_id},
+                "permission": permission,
+            },
+            safe_to_retry=True,
+        )
+
+    # --- sessions ------------------------------------------------------------
+
+    def create_session(
+        self,
+        store: str,
+        actor_id: str,
+        session_id: Optional[str] = None,
+        parent_session_id: Optional[str] = None,
+        metadata: Optional[dict] = None,
+    ) -> models.Session:
+        body = _body(actor_id=actor_id, parent_session_id=parent_session_id, metadata=metadata)
+        return _as(
+            models.Session,
+            self._do(
+                "POST",
+                f"{_BASE}/session-stores/{store}/sessions",
+                query=_query(session_id=session_id),
+                body=body,
+            ),
+        )
+
+    def list_sessions(
+        self,
+        store: str,
+        filter: Optional[str] = None,
+        order_by: Optional[str] = None,
+        page_size: Optional[int] = None,
+        page_token: Optional[str] = None,
+    ) -> models.SessionList:
+        return _as(
+            models.SessionList,
+            self._do(
+                "GET",
+                f"{_BASE}/session-stores/{store}/sessions",
+                query=_query(
+                    filter=filter, order_by=order_by, page_size=page_size, page_token=page_token
+                ),
+            ),
+        )
+
+    def get_session(self, session_id: str, store: str) -> models.Session:
+        path = f"{_BASE}/session-stores/{store}/sessions/{session_id}"
+        return _as(models.Session, self._do("GET", path))
+
+    def update_session(self, store: str, session_id: str, metadata: dict) -> models.Session:
+        return _as(
+            models.Session,
+            self._do(
+                "PATCH",
+                f"{_BASE}/session-stores/{store}/sessions/{session_id}",
+                query={"update_mask": "metadata"},
+                body=_body(metadata=metadata),
+            ),
+        )
+
+    def delete_session(self, store: str, session_id: str, force: bool = False) -> dict:
+        return self._do(
+            "DELETE",
+            f"{_BASE}/session-stores/{store}/sessions/{session_id}",
+            query={"force": True} if force else None,
+        )
+
+    def fork_session(
+        self,
+        store: str,
+        source_session_id: str,
+        actor_id: str,
+        up_to_item_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        metadata: Optional[dict] = None,
+    ) -> models.Session:
+        body = _body(
+            source_session_id=source_session_id,
+            actor_id=actor_id,
+            up_to_item_id=up_to_item_id,
+            session_id=session_id,
+            metadata=metadata,
+        )
+        return _as(
+            models.Session,
+            self._do("POST", f"{_BASE}/session-stores/{store}/sessions:fork", body=body),
+        )
+
+    # --- session items -------------------------------------------------------
+
+    def list_session_items(
+        self,
+        store: str,
+        session_id: str,
+        order_by: Optional[str] = None,
+        page_size: Optional[int] = None,
+        page_token: Optional[str] = None,
+    ) -> models.SessionItemList:
+        return _as(
+            models.SessionItemList,
+            self._do(
+                "GET",
+                f"{_BASE}/session-stores/{store}/sessions/{session_id}/items",
+                query=_query(order_by=order_by, page_size=page_size, page_token=page_token),
+            ),
+        )
+
+    def append_session_items(
+        self, store: str, session_id: str, items: list[dict]
+    ) -> models.SessionItemList:
+        body = {"items": [{"data": item} for item in items]}
+        return _as(
+            models.SessionItemList,
+            self._do(
+                "POST",
+                f"{_BASE}/session-stores/{store}/sessions/{session_id}/items:append",
+                body=body,
+            ),
+        )
+
+    def pop_session_item(self, store: str, session_id: str) -> models.PoppedSessionItem:
+        return _as(
+            models.PoppedSessionItem,
+            self._do(
+                "POST", f"{_BASE}/session-stores/{store}/sessions/{session_id}/items:pop", body={}
+            ),
+        )
+
+    def clear_session_items(self, store: str, session_id: str) -> dict:
+        return self._do(
+            "POST", f"{_BASE}/session-stores/{store}/sessions/{session_id}/items:clear", body={}
+        )
+
+    def extract_memories(
+        self,
+        store: str,
+        session_id: str,
+        memory_store: str,
+        *,
+        instructions: Optional[str] = None,
+        dry_run: bool = False,
+    ) -> models.ExtractMemoriesResponse:
+        """Synchronously distill a session's transcript into memory entries.
+
+        Writes the entries to `memory_store` and returns them; pass `dry_run=True` to return the
+        extracted entries without persisting them.
+        """
+        body = _body(
+            memory_store=memory_store_path(memory_store),
+            instructions=instructions,
+            dry_run=dry_run or None,
+        )
+        return _as(
+            models.ExtractMemoriesResponse,
+            self._do(
+                "POST",
+                f"{_BASE}/session-stores/{store}/sessions/{session_id}/extractions",
+                body=body,
+            ),
+        )

databricks_agentkit/_pagination.py

@@ -0,0 +1,25 @@
+"""Shared pagination validation for typed Agent Bricks resources."""
+
+from typing import Optional
+
+_MAX_PAGE_SIZE = 100
+
+
+def _validate_bound(value: Optional[int], name: str) -> None:
+    if value is None:
+        return
+    # Reject non-integers (and bool, which is an int subclass) up front: without this a str
+    # raised a cryptic ``'<=' not supported between 'int' and 'str'`` TypeError, and a bool/float
+    # slipped through the range check (True == 1, 1.5 in [1, 100]) to reach the API as garbage.
+    if not isinstance(value, int) or isinstance(value, bool):
+        raise ValueError(f"{name} must be an integer between 1 and {_MAX_PAGE_SIZE}")
+    if not 1 <= value <= _MAX_PAGE_SIZE:
+        raise ValueError(f"{name} must be between 1 and {_MAX_PAGE_SIZE}")
+
+
+def validate_page_size(page_size: Optional[int]) -> None:
+    _validate_bound(page_size, "page_size")
+
+
+def validate_limit(limit: Optional[int]) -> None:
+    _validate_bound(limit, "limit")