databricks-agentbricks 0.2.0__py3-none-any.whl

databricks_agentkit/runtime/model_services.py

@@ -0,0 +1,90 @@
+"""Discover the Unity Catalog AI Gateway chat model services available in a workspace.
+
+The AI Gateway exposes Databricks-managed models as Unity Catalog *model services* in the
+``system.ai`` schema (e.g. ``system.ai.claude-sonnet-4-5``), queryable through an OpenAI-compatible
+endpoint at ``<host>/ai-gateway/mlflow/v1``. This lists the chat-capable ones so the demo UI's model
+picker can offer them; the agent then calls the chosen one with ``use_ai_gateway=True``.
+
+Kept framework-neutral (no agent SDK, no ``mlflow``) so it sits with the other neutral runtime
+helpers and can be listed from either template's ``runtime/ui.py``.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import TYPE_CHECKING, Any, cast
+
+if TYPE_CHECKING:
+    from databricks.sdk import WorkspaceClient
+
+# The Unity Catalog REST route that lists model services under a schema.
+_MODEL_SERVICES_PATH = "/api/2.1/unity-catalog/model-services"
+# The AI Gateway's Databricks-managed models live in this UC schema.
+_SYSTEM_AI_SCHEMA = "system.ai"
+# The list API rejects page_size above 100 (InvalidParameterValue), so use its maximum.
+_PAGE_SIZE = 100
+# Resource-name prefix the API returns, e.g. "model-services/system.ai.claude-sonnet-4-5".
+_NAME_PREFIX = "model-services/"
+# The list call can intermittently fail; retry a few times so a transient blip doesn't blank the
+# picker. A non-transient error (e.g. no permission) simply fails every attempt, then surfaces.
+_LIST_ATTEMPTS = 3
+_RETRY_DELAY_S = 0.5
+
+
+def _is_chat_capable(supported_api_types: Any) -> bool:
+    """True if a model service speaks a chat/completions API (i.e. not embeddings-only).
+
+    ``supported_api_types`` looks like ``["openai/v1/chat/completions"]``. Be lenient when it's
+    absent (offer the model), but drop services that only advertise embeddings.
+    """
+    types = [str(t).lower() for t in (supported_api_types or [])]
+    if not types:
+        return True
+    return any(("chat" in t or "completions" in t or "responses" in t) for t in types)
+
+
+def _list_page(client: WorkspaceClient, query: dict[str, Any]) -> Any:
+    """Fetch one page of the model-services list, retrying transient failures."""
+    last_error: Exception | None = None
+    for attempt in range(_LIST_ATTEMPTS):
+        try:
+            return client.api_client.do("GET", _MODEL_SERVICES_PATH, query=query)
+        except Exception as exc:  # noqa: BLE001 - retry any list failure, then surface the last one
+            last_error = exc
+            if attempt < _LIST_ATTEMPTS - 1:
+                time.sleep(_RETRY_DELAY_S)
+    raise last_error if last_error else RuntimeError("model-services list returned nothing")
+
+
+def list_ai_gateway_model_services(client: WorkspaceClient) -> list[str]:
+    """Return the names of chat-capable ``system.ai`` AI Gateway model services, sorted.
+
+    Each name is a Unity Catalog model-service path like ``system.ai.claude-sonnet-4-5`` — exactly
+    the string the OpenAI-compatible gateway (``use_ai_gateway=True``) expects as its ``model``.
+    Pages through the model-services list API, retrying transient list failures per page.
+
+    Raises whatever the underlying request raises after retries (e.g. a permission error); callers
+    that want a graceful fallback to just the agent's default model should catch it.
+    """
+    names: list[str] = []
+    page_token: str | None = None
+    while True:
+        query: dict[str, Any] = {"parent": f"schemas/{_SYSTEM_AI_SCHEMA}", "page_size": _PAGE_SIZE}
+        if page_token:
+            query["page_token"] = page_token
+        raw = _list_page(client, query)
+        if not isinstance(raw, dict):
+            break
+        response = cast("dict[str, Any]", raw)
+        for service in response.get("model_services") or []:
+            if not isinstance(service, dict):
+                continue
+            name = str(service.get("name") or "")
+            if name.startswith(_NAME_PREFIX):
+                name = name[len(_NAME_PREFIX) :]
+            if name and _is_chat_capable(service.get("supported_api_types")):
+                names.append(name)
+        page_token = response.get("next_page_token") or None
+        if not page_token:
+            break
+    return sorted(names)

databricks_agentkit/runtime/runtime.py

@@ -0,0 +1,273 @@
+"""Transport-neutral facade for Agent Bricks invocation execution."""
+
+from __future__ import annotations
+
+import asyncio
+import copy
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+from databricks_agentkit.runtime.execution import (
+    AttemptExecution,
+    InvocationExecutor,
+    LocalInvocationExecutor,
+    copy_json_value,
+)
+from databricks_agentkit.runtime.store import (
+    InMemoryRuntimeStore,
+    RuntimeStore,
+    runtime_store_from_environment,
+)
+from databricks_agentkit.runtime.types import (
+    Invocation,
+    InvocationEvent,
+    InvocationExecutorFn,
+    InvocationFailedError,
+    InvocationNotFoundError,
+    InvocationStatus,
+    JsonValue,
+)
+
+if TYPE_CHECKING:
+    from databricks_agentkit.runtime.durability.store import DurableRuntimeStore
+
+
+class Runtime:
+    """Coordinate one Agent Bricks process's invocation lifecycle.
+
+    A Runtime Store owns invocation state: requests, status, events, and results. An invocation
+    executor turns queued state into agent attempts. This facade accepts idempotent requests,
+    delegates scheduling to the executor, and polls the store for foreground results or event replay.
+    Use :meth:`local` for process-local execution, :meth:`durable` for a durable store,
+    :meth:`from_store` for an explicit store, or :meth:`from_environment` when Agent Bricks configures the
+    store through the process environment.
+    """
+
+    def __init__(
+        self,
+        *,
+        runtime_store: RuntimeStore,
+        executor: InvocationExecutor,
+        poll_seconds: float = 1.0,
+    ) -> None:
+        if poll_seconds <= 0:
+            raise ValueError("poll_seconds must be positive")
+        self.runtime_store = runtime_store
+        self.executor = executor
+        self.poll_seconds = poll_seconds
+        self._started = False
+
+    @classmethod
+    def local(
+        cls,
+        execute_fn: InvocationExecutorFn,
+        *,
+        runtime_store: RuntimeStore | None = None,
+        poll_seconds: float = 1.0,
+    ) -> Runtime:
+        """Construct a Runtime that schedules work only in the current process.
+
+        The default :class:`InMemoryRuntimeStore` supports foreground, background, streaming,
+        polling, and event replay until the process exits. It does not heartbeat or recover work.
+        """
+        store = InMemoryRuntimeStore() if runtime_store is None else runtime_store
+        from databricks_agentkit.runtime.durability.store import DurableRuntimeStore
+
+        if isinstance(store, DurableRuntimeStore):
+            raise TypeError("Runtime.local requires a non-durable RuntimeStore")
+        execution = AttemptExecution(execute_fn, runtime_store=store)
+        return cls(
+            runtime_store=store,
+            executor=LocalInvocationExecutor(execution, runtime_store=store),
+            poll_seconds=poll_seconds,
+        )
+
+    @classmethod
+    def durable(
+        cls,
+        execute_fn: InvocationExecutorFn,
+        *,
+        runtime_store: DurableRuntimeStore,
+        recovery_enabled: Callable[[], bool] | None = None,
+        heartbeat_seconds: float = 3.0,
+        stale_seconds: float = 10.0,
+        scan_seconds: float = 3.0,
+        poll_seconds: float = 1.0,
+    ) -> Runtime:
+        """Construct a Runtime that heartbeats and can recover work from a durable store.
+
+        A durable runtime requires a :class:`DurableRuntimeStore`. Its results and events persist
+        across process restarts. ``recovery_enabled`` controls only stale-work recovery; it does not
+        change whether normal invocation state is persisted.
+        """
+        from databricks_agentkit.runtime.durability.execution import DurableInvocationExecutor
+        from databricks_agentkit.runtime.durability.store import DurableRuntimeStore
+
+        if not isinstance(runtime_store, DurableRuntimeStore):
+            raise TypeError("Runtime.durable requires a DurableRuntimeStore")
+        execution = AttemptExecution(execute_fn, runtime_store=runtime_store)
+        return cls(
+            runtime_store=runtime_store,
+            executor=DurableInvocationExecutor(
+                execution,
+                runtime_store=runtime_store,
+                recovery_enabled=recovery_enabled or (lambda: True),
+                heartbeat_seconds=heartbeat_seconds,
+                stale_seconds=stale_seconds,
+                scan_seconds=scan_seconds,
+            ),
+            poll_seconds=poll_seconds,
+        )
+
+    @classmethod
+    def from_store(
+        cls,
+        execute_fn: InvocationExecutorFn,
+        *,
+        runtime_store: RuntimeStore,
+        recovery_enabled: Callable[[], bool] | None = None,
+        heartbeat_seconds: float = 3.0,
+        stale_seconds: float = 10.0,
+        scan_seconds: float = 3.0,
+        poll_seconds: float = 1.0,
+    ) -> Runtime:
+        """Construct the local or durable Runtime required by the selected Runtime Store."""
+        from databricks_agentkit.runtime.durability.store import DurableRuntimeStore
+
+        if isinstance(runtime_store, DurableRuntimeStore):
+            return cls.durable(
+                execute_fn,
+                runtime_store=runtime_store,
+                recovery_enabled=recovery_enabled,
+                heartbeat_seconds=heartbeat_seconds,
+                stale_seconds=stale_seconds,
+                scan_seconds=scan_seconds,
+                poll_seconds=poll_seconds,
+            )
+        return cls.local(execute_fn, runtime_store=runtime_store, poll_seconds=poll_seconds)
+
+    @classmethod
+    def from_environment(
+        cls,
+        execute_fn: InvocationExecutorFn,
+        *,
+        recovery_enabled: Callable[[], bool] | None = None,
+        heartbeat_seconds: float = 3.0,
+        stale_seconds: float = 10.0,
+        scan_seconds: float = 3.0,
+        poll_seconds: float = 1.0,
+    ) -> Runtime:
+        """Construct the Runtime selected by the Agent Bricks process environment."""
+        return cls.from_store(
+            execute_fn,
+            runtime_store=runtime_store_from_environment(),
+            recovery_enabled=recovery_enabled,
+            heartbeat_seconds=heartbeat_seconds,
+            stale_seconds=stale_seconds,
+            scan_seconds=scan_seconds,
+            poll_seconds=poll_seconds,
+        )
+
+    @property
+    def is_durable(self) -> bool:
+        """Whether this Runtime has heartbeat and recovery capabilities."""
+        from databricks_agentkit.runtime.durability.store import DurableRuntimeStore
+
+        return isinstance(self.runtime_store, DurableRuntimeStore)
+
+    async def start(self) -> None:
+        """Initialize the Runtime Store and invocation executor."""
+        if self._started:
+            return
+        await self.runtime_store.initialize()
+        try:
+            await self.executor.start()
+        except Exception:
+            await self.runtime_store.close()
+            raise
+        self._started = True
+
+    async def stop(self) -> None:
+        """Stop local work and close the Runtime Store."""
+        if not self._started:
+            return
+        await self.executor.stop()
+        self._started = False
+        await self.runtime_store.close()
+
+    async def submit(self, invocation_id: str, request: JsonValue) -> Invocation:
+        """Accept an idempotent invocation and schedule it on this worker.
+
+        The Runtime Store makes ``invocation_id`` the idempotency key. Retrying the same request
+        returns its existing state; a different request for the same ID raises
+        :class:`InvocationConflictError`.
+        """
+        self._require_started()
+        if not invocation_id:
+            raise ValueError("invocation_id must not be empty")
+        state = await self.runtime_store.accept(
+            invocation_id,
+            copy_json_value(request, "request"),
+        )
+        self.executor.ensure_scheduled(state)
+        return state
+
+    async def invoke(
+        self,
+        invocation_id: str,
+        request: JsonValue,
+        *,
+        timeout: float | None = None,
+    ) -> JsonValue:
+        """Accept an invocation and wait for its terminal result."""
+        await self.submit(invocation_id, request)
+        return await self.wait(invocation_id, timeout=timeout)
+
+    async def get_invocation(self, invocation_id: str) -> Invocation | None:
+        """Return current state and ensure queued first-attempt work is scheduled locally."""
+        self._require_started()
+        state = await self.runtime_store.get(invocation_id)
+        if state is not None:
+            self.executor.ensure_scheduled(state)
+        return state
+
+    async def wait(
+        self,
+        invocation_id: str,
+        *,
+        timeout: float | None = None,
+    ) -> JsonValue:
+        """Poll until an invocation completes, fails, or reaches the optional timeout."""
+        self._require_started()
+
+        async def poll() -> JsonValue:
+            while True:
+                state = await self.get_invocation(invocation_id)
+                if state is None:
+                    raise InvocationNotFoundError(invocation_id)
+                if state.status == InvocationStatus.COMPLETED:
+                    return copy.deepcopy(state.response)
+                if state.status == InvocationStatus.FAILED:
+                    raise InvocationFailedError(invocation_id)
+                await asyncio.sleep(self.poll_seconds)
+
+        if timeout is None:
+            return await poll()
+        try:
+            return await asyncio.wait_for(poll(), timeout=timeout)
+        except asyncio.TimeoutError as exc:
+            raise TimeoutError from exc
+
+    async def get_events(
+        self,
+        invocation_id: str,
+        *,
+        after_sequence: int | None = None,
+    ) -> list[InvocationEvent]:
+        """Return persisted invocation events after an optional exclusive replay cursor."""
+        self._require_started()
+        return await self.runtime_store.events(invocation_id, after_sequence)
+
+    def _require_started(self) -> None:
+        if not self._started:
+            raise RuntimeError("Runtime.start() must be called first")

databricks_agentkit/runtime/session_store_client.py

@@ -0,0 +1,120 @@
+"""Minimal Databricks Session Store REST client (vendored).
+
+Wraps the managed Session Store API (``/api/2.0/agents/session-stores``) — the SDK-agnostic durable
+store for agent session history — with just the calls the LangGraph saver needs: resolve a store,
+get-or-create a session, and append / list / clear its ordered items. Auth/host come from a
+``databricks.sdk.WorkspaceClient`` (default credentials).
+
+This is a trimmed stand-in for the first-party ``databricks_agent_client.SessionStoreClient``; when
+that package is published, import it instead and delete this file. Item ``data`` is opaque JSON, so
+callers store whatever shape they like (the saver stores checkpoint fragments).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Iterator, Optional, Sequence
+
+from databricks.sdk import WorkspaceClient
+
+from databricks_agentkit.runtime.workspace import workspace_client as default_workspace_client
+
+_API_ROOT = "/api/2.0/agents"
+
+
+@dataclass(frozen=True)
+class Session:
+    """A durable interaction within a Session Store, addressed by (store, session_id)."""
+
+    session_store_name: str
+    session_id: str
+    actor_id: str
+    metadata: dict[str, str] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class SessionItem:
+    """One opaque, ordered history entry. ``data`` is arbitrary JSON."""
+
+    item_id: str
+    data: Any
+
+
+class SessionStoreClient:
+    """Thin REST client over the managed Session Store API."""
+
+    def __init__(self, workspace_client: Optional[WorkspaceClient] = None) -> None:
+        self._api = (workspace_client or default_workspace_client()).api_client
+        self._store_name: Optional[str] = None
+
+    def set_session_store(self, session_store_name: str) -> "SessionStoreClient":
+        """Bind this client to a store; returns self so it doubles as the store handle."""
+        if not session_store_name:
+            raise ValueError("session_store_name is required")
+        self._store_name = session_store_name
+        return self
+
+    def get_session(self, *, session_id: str) -> Session:
+        resp: dict[str, Any] = self._api.do("GET", f"{self._sessions_path()}/{session_id}")  # type: ignore[assignment]
+        return self._session(resp)
+
+    def create_session(
+        self,
+        *,
+        actor_id: str,
+        session_id: Optional[str] = None,
+        metadata: Optional[dict[str, str]] = None,
+    ) -> Session:
+        body: dict[str, Any] = {"actor_id": actor_id}
+        if metadata:
+            body["metadata"] = metadata
+        query = {"session_id": session_id} if session_id else None
+        resp: dict[str, Any] = self._api.do("POST", self._sessions_path(), query=query, body=body)  # type: ignore[assignment]
+        return self._session(resp)
+
+    def append_items(self, session: Session, *, items: Sequence[Any]) -> None:
+        if not items:
+            raise ValueError("at least one item is required")
+        self._api.do(
+            "POST",
+            f"{self._items_path(session)}:append",
+            body={"items": [{"data": item} for item in items]},
+        )
+
+    def list_items(
+        self, session: Session, *, order_by: Optional[str] = None
+    ) -> Iterator[SessionItem]:
+        """Yield the session's items (paginating transparently)."""
+        page_token: Optional[str] = None
+        while True:
+            query = {k: v for k, v in {"order_by": order_by, "page_token": page_token}.items() if v}
+            resp: dict[str, Any] = self._api.do(
+                "GET", self._items_path(session), query=query or None
+            )  # type: ignore[assignment]
+            for item in resp.get("session_items", []):
+                if "data" in item:
+                    yield SessionItem(item_id=item.get("item_id", ""), data=item["data"])
+            page_token = resp.get("next_page_token")
+            if not page_token:
+                return
+
+    def clear_items(self, session: Session) -> None:
+        self._api.do("POST", f"{self._items_path(session)}:clear", body={})
+
+    # ----- internals ----------------------------------------------------------
+
+    def _sessions_path(self) -> str:
+        if not self._store_name:
+            raise ValueError("call set_session_store() before using the client")
+        return f"{_API_ROOT}/session-stores/{self._store_name}/sessions"
+
+    def _items_path(self, session: Session) -> str:
+        return f"{self._sessions_path()}/{session.session_id}/items"
+
+    def _session(self, resp: dict[str, Any]) -> Session:
+        return Session(
+            session_store_name=resp.get("session_store_name", self._store_name or ""),
+            session_id=resp["session_id"],
+            actor_id=resp.get("actor_id", ""),
+            metadata=dict(resp.get("metadata", {})),
+        )