databricks-agentbricks 0.2.0__py3-none-any.whl

databricks_agentkit/runtime/types.py

@@ -0,0 +1,118 @@
+"""Contracts shared by Agent Bricks Runtime, stores, and agent applications.
+
+The runtime returns :class:`Invocation` snapshots and ordered :class:`InvocationEvent` records.
+An executor receives an :class:`InvocationAttemptContext` for attempt fencing and event emission.
+``DurableAgentServer`` adapts that lower-level context into :class:`InvocationContext` for functions
+registered with ``@app.invoke`` and ``@app.recover``.
+
+All request, response, and event payloads use the recursive ``JsonValue`` / ``JsonObject`` aliases,
+so Runtime Store values can be stored consistently by in-memory and Lakebase implementations.
+"""
+
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import TYPE_CHECKING, TypeAlias
+
+if TYPE_CHECKING:
+    from databricks_agentkit.runtime.auth import RequestAuthContext
+
+JsonValue: TypeAlias = None | bool | int | float | str | list["JsonValue"] | dict[str, "JsonValue"]
+JsonObject = dict[str, JsonValue]
+InvocationEventEmitter = Callable[[JsonObject], Awaitable[int]]
+
+
+class InvocationStatus(str, Enum):
+    """Lifecycle states stored by Agent Bricks Runtime."""
+
+    QUEUED = "QUEUED"
+    ACTIVE = "ACTIVE"
+    COMPLETED = "COMPLETED"
+    FAILED = "FAILED"
+
+
+@dataclass(frozen=True)
+class Invocation:
+    """A Runtime Store snapshot for one idempotent invocation and its owning attempt."""
+
+    invocation_id: str
+    status: InvocationStatus
+    attempt: int
+    request: JsonValue
+    response: JsonValue
+
+    @property
+    def is_terminal(self) -> bool:
+        """Whether this invocation can no longer transition or emit events."""
+        return self.status in {
+            InvocationStatus.COMPLETED,
+            InvocationStatus.FAILED,
+        }
+
+
+@dataclass(frozen=True)
+class InvocationEvent:
+    """One persisted event emitted by an invocation attempt."""
+
+    sequence_number: int
+    invocation_id: str
+    attempt: int
+    event: JsonObject
+
+
+@dataclass(frozen=True)
+class InvocationAttemptContext:
+    """Attempt metadata and event emission passed to the runtime's executor function."""
+
+    invocation_id: str
+    attempt: int
+    _emit: InvocationEventEmitter | None = field(default=None, repr=False, compare=False)
+
+    @property
+    def is_recovery(self) -> bool:
+        """Whether this is a replacement attempt after an earlier worker stopped heartbeating."""
+        return self.attempt > 1
+
+    async def emit(self, event: JsonObject) -> int:
+        """Persist an ordered event and return its replay cursor."""
+        if self._emit is None:
+            raise RuntimeError("event emission is not available for this invocation context")
+        return await self._emit(event)
+
+
+InvocationExecutorFn = Callable[[JsonValue, InvocationAttemptContext], Awaitable[JsonValue]]
+
+
+@dataclass(frozen=True)
+class InvocationContext:
+    """Invocation/session metadata and event emission for a decorated agent function."""
+
+    invocation_id: str
+    session_id: str
+    attempt: int
+    _attempt_context: InvocationAttemptContext = field(repr=False, compare=False)
+    request_auth: "RequestAuthContext | None" = field(default=None, repr=False, compare=False)
+
+    @property
+    def is_recovery(self) -> bool:
+        """Whether ``@app.recover`` is handling a replacement attempt."""
+        return self.attempt > 1
+
+    async def emit(self, event: JsonObject) -> int:
+        """Persist an ordered application event and return its replay cursor."""
+        return await self._attempt_context.emit(event)
+
+
+InvocationHook = Callable[[JsonValue, InvocationContext], Awaitable[JsonValue]]
+
+
+class InvocationConflictError(ValueError):
+    """Raised when an invocation ID is reused with a different request."""
+
+
+class InvocationNotFoundError(LookupError):
+    """Raised when waiting for an unknown invocation ID."""
+
+
+class InvocationFailedError(RuntimeError):
+    """Raised when an invocation reaches the failed state."""

databricks_agentkit/runtime/workspace.py

@@ -0,0 +1,27 @@
+"""Construct Databricks SDK clients with workspace routing when required."""
+
+from __future__ import annotations
+
+import os
+
+from databricks.sdk import WorkspaceClient
+
+
+def workspace_headers() -> dict[str, str]:
+    """Return headers required to route an account-host request to its workspace."""
+    workspace_id = os.getenv("DATABRICKS_WORKSPACE_ID", "").strip()
+    return {"X-Databricks-Org-Id": workspace_id} if workspace_id else {}
+
+
+def workspace_client() -> WorkspaceClient:
+    """Return the environment-authenticated client for the active workspace.
+
+    ``databricks apps run-local`` can authenticate through an account-level vanity host while
+    exposing the target workspace through ``DATABRICKS_WORKSPACE_ID``. The SDK needs the same
+    routing header as the Databricks CLI for those profiles. Ordinary workspace hosts and deployed
+    Apps continue to use the SDK's default authentication chain.
+    """
+    headers = workspace_headers()
+    if not headers:
+        return WorkspaceClient()
+    return WorkspaceClient(custom_headers=headers)

databricks_agentkit/session_store.py

@@ -0,0 +1,434 @@
+"""Typed session-store resources and operations."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator, Sequence
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, List, Optional
+
+from databricks_agentkit._pagination import validate_page_size
+from databricks_agentkit.timefmt import parse_timestamp
+
+if TYPE_CHECKING:
+    from databricks_agentkit._api_client import _AgentBricksApiClient
+
+
+@dataclass(frozen=True)
+class SessionItem:
+    item_id: str
+    data: Any
+    create_time: Optional[datetime] = None
+
+
+@dataclass(frozen=True, kw_only=True)
+class ExtractedMemory:
+    """A memory entry produced by `Session.extract_memories` (a read-only result view)."""
+
+    name: str
+    actor_id: Optional[str] = None
+    path: Optional[str] = None
+    content: Optional[str] = None
+    description: Optional[str] = None
+    session_id: Optional[str] = None
+    source_type: Optional[str] = None
+    create_time: Optional[datetime] = None
+    update_time: Optional[datetime] = None
+
+
+@dataclass(frozen=True, kw_only=True)
+class Session:
+    store_name: str
+    session_id: str
+    actor_id: str
+    parent_session_id: Optional[str] = None
+    root_session_id: Optional[str] = None
+    metadata: dict[str, str] = field(default_factory=dict)
+    create_time: Optional[datetime] = None
+    update_time: Optional[datetime] = None
+    last_activity_time: Optional[datetime] = None
+    _client: SessionStores = field(repr=False, compare=False)
+
+    def update(self, *, metadata: dict[str, str]) -> Session:
+        return self._client._update_session(self, metadata=metadata)
+
+    def delete(self, *, force: bool = False) -> None:
+        self._client._delete_session(self, force=force)
+
+    def fork(
+        self,
+        *,
+        actor_id: str,
+        up_to_item_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        metadata: Optional[dict[str, str]] = None,
+    ) -> Session:
+        return self._client._fork_session(
+            self,
+            actor_id=actor_id,
+            up_to_item_id=up_to_item_id,
+            session_id=session_id,
+            metadata=metadata,
+        )
+
+    def list_items(
+        self,
+        *,
+        page_size: Optional[int] = None,
+        order_by: Optional[str] = None,
+    ) -> Iterator[SessionItem]:
+        return self._client._list_items(
+            self,
+            page_size=page_size,
+            order_by=order_by,
+        )
+
+    def append_items(self, items: Sequence[Any]) -> List[SessionItem]:
+        return self._client._append_items(self, items=items)
+
+    def pop_item(self) -> Optional[SessionItem]:
+        return self._client._pop_item(self)
+
+    def clear_items(self) -> None:
+        self._client._clear_items(self)
+
+    def extract_memories(
+        self,
+        *,
+        memory_store: str,
+        instructions: Optional[str] = None,
+        dry_run: bool = False,
+    ) -> List[ExtractedMemory]:
+        """Synchronously distill this 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.
+        """
+        return self._client._extract_memories(
+            self, memory_store=memory_store, instructions=instructions, dry_run=dry_run
+        )
+
+
+@dataclass(frozen=True, kw_only=True)
+class SessionStore:
+    name: str
+    session_store_id: Optional[str] = None
+    creator_user_id: Optional[str] = None
+    create_time: Optional[datetime] = None
+    update_time: Optional[datetime] = None
+    description: Optional[str] = None
+    metadata: dict[str, str] = field(default_factory=dict)
+    _client: SessionStores = field(repr=False, compare=False)
+
+    def update(
+        self,
+        *,
+        description: Optional[str] = None,
+        metadata: Optional[dict[str, str]] = None,
+    ) -> SessionStore:
+        return self._client._update_store(
+            self,
+            description=description,
+            metadata=metadata,
+        )
+
+    def delete(self) -> None:
+        self._client._delete_store(self)
+
+    def add(
+        self,
+        *,
+        actor_id: str,
+        session_id: Optional[str] = None,
+        parent_session_id: Optional[str] = None,
+        metadata: Optional[dict[str, str]] = None,
+    ) -> Session:
+        return self._client._add_session(
+            self,
+            actor_id=actor_id,
+            session_id=session_id,
+            parent_session_id=parent_session_id,
+            metadata=metadata,
+        )
+
+    def list(
+        self,
+        *,
+        page_size: Optional[int] = None,
+        filter: Optional[str] = None,
+        order_by: Optional[str] = None,
+    ) -> Iterator[Session]:
+        return self._client._list_sessions(
+            self,
+            page_size=page_size,
+            filter=filter,
+            order_by=order_by,
+        )
+
+    def get(self, session_id: str) -> Session:
+        return self._client._get_session(self, session_id=session_id)
+
+    def grant_permission(self, principal_id: str, *, permission: str = "WRITE") -> None:
+        """Grant a service principal READ/WRITE access to this session store.
+
+        ``principal_id`` is the service principal's application (client) id. The grant is applied
+        server-side, so the caller needs neither store ownership nor Lakebase MANAGE.
+        """
+        self._client._grant_permission(self, principal_id, permission=permission)
+
+
+class SessionStores:
+    def __init__(self, api: _AgentBricksApiClient):
+        self._api = api
+
+    def create(
+        self,
+        name: str,
+        *,
+        description: Optional[str] = None,
+        metadata: Optional[dict[str, str]] = None,
+    ) -> SessionStore:
+        response = self._api.create_session_store(
+            name,
+            description,
+            metadata,
+        )
+        return self._store_from_response(response)
+
+    def list(self, *, page_size: Optional[int] = None) -> Iterator[SessionStore]:
+        validate_page_size(page_size)
+        page_token = None
+        while True:
+            response = self._api.list_session_stores(
+                page_size=page_size,
+                page_token=page_token,
+            )
+            for store in response.get("session_stores", []):
+                yield self._store_from_response(store)
+            page_token = response.get("next_page_token")
+            if not page_token:
+                return
+
+    def get(self, name: str) -> SessionStore:
+        return self._store_from_response(self._api.get_session_store(name))
+
+    def _update_store(
+        self,
+        store: SessionStore,
+        *,
+        description: Optional[str] = None,
+        metadata: Optional[dict[str, str]] = None,
+    ) -> SessionStore:
+        if description is None and metadata is None:
+            raise ValueError("at least one of description and metadata is required")
+        response = self._api.update_session_store(
+            store.name,
+            description=description,
+            metadata=metadata,
+        )
+        return self._store_from_response(response)
+
+    def _delete_store(self, store: SessionStore) -> None:
+        self._api.delete_session_store(store.name)
+
+    def _grant_permission(self, store: SessionStore, principal_id: str, *, permission: str) -> None:
+        self._api.grant_session_store_permission(store.name, principal_id, permission=permission)
+
+    def _add_session(
+        self,
+        store: SessionStore,
+        *,
+        actor_id: str,
+        session_id: Optional[str] = None,
+        parent_session_id: Optional[str] = None,
+        metadata: Optional[dict[str, str]] = None,
+    ) -> Session:
+        response = self._api.create_session(
+            store.name,
+            actor_id,
+            session_id=session_id,
+            parent_session_id=parent_session_id,
+            metadata=metadata,
+        )
+        return self._session_from_response(response, store.name)
+
+    def _list_sessions(
+        self,
+        store: SessionStore,
+        *,
+        page_size: Optional[int] = None,
+        filter: Optional[str] = None,
+        order_by: Optional[str] = None,
+    ) -> Iterator[Session]:
+        validate_page_size(page_size)
+        if order_by is None:
+            order_by = "create_time desc"
+        page_token = None
+        while True:
+            response = self._api.list_sessions(
+                store.name,
+                filter=filter,
+                order_by=order_by,
+                page_size=page_size,
+                page_token=page_token,
+            )
+            for session in response.get("sessions", []):
+                yield self._session_from_response(session, store.name)
+            page_token = response.get("next_page_token")
+            if not page_token:
+                return
+
+    def _get_session(
+        self,
+        store: SessionStore,
+        *,
+        session_id: str,
+    ) -> Session:
+        response = self._api.get_session(session_id, store.name)
+        return self._session_from_response(response, store.name)
+
+    def _update_session(self, session: Session, *, metadata: dict[str, str]) -> Session:
+        response = self._api.update_session(
+            session.store_name,
+            session.session_id,
+            metadata,
+        )
+        return self._session_from_response(response, session.store_name)
+
+    def _delete_session(self, session: Session, *, force: bool = False) -> None:
+        self._api.delete_session(session.store_name, session.session_id, force=force)
+
+    def _fork_session(
+        self,
+        session: Session,
+        *,
+        actor_id: str,
+        up_to_item_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        metadata: Optional[dict[str, str]] = None,
+    ) -> Session:
+        response = self._api.fork_session(
+            session.store_name,
+            session.session_id,
+            actor_id,
+            up_to_item_id=up_to_item_id,
+            session_id=session_id,
+            metadata=metadata,
+        )
+        return self._session_from_response(
+            response.get("session", response),
+            session.store_name,
+        )
+
+    def _list_items(
+        self,
+        session: Session,
+        *,
+        page_size: Optional[int] = None,
+        order_by: Optional[str] = None,
+    ) -> Iterator[SessionItem]:
+        validate_page_size(page_size)
+        page_token = None
+        while True:
+            response = self._api.list_session_items(
+                session.store_name,
+                session.session_id,
+                order_by=order_by,
+                page_size=page_size,
+                page_token=page_token,
+            )
+            for item in response.get("session_items", []):
+                yield self._item_from_response(item)
+            page_token = response.get("next_page_token")
+            if not page_token:
+                return
+
+    def _append_items(self, session: Session, *, items: Sequence[Any]) -> List[SessionItem]:
+        if not items:
+            raise ValueError("at least one item is required")
+        response = self._api.append_session_items(
+            session.store_name,
+            session.session_id,
+            list(items),
+        )
+        return [self._item_from_response(item) for item in response.get("session_items", [])]
+
+    def _pop_item(self, session: Session) -> Optional[SessionItem]:
+        response = self._api.pop_session_item(session.store_name, session.session_id)
+        item = response.get("item")
+        return self._item_from_response(item) if item is not None else None
+
+    def _clear_items(self, session: Session) -> None:
+        self._api.clear_session_items(session.store_name, session.session_id)
+
+    def _extract_memories(
+        self,
+        session: Session,
+        *,
+        memory_store: str,
+        instructions: Optional[str],
+        dry_run: bool,
+    ) -> List[ExtractedMemory]:
+        response = self._api.extract_memories(
+            session.store_name,
+            session.session_id,
+            memory_store,
+            instructions=instructions,
+            dry_run=dry_run,
+        )
+        return [self._extracted_from_response(e) for e in response.get("entries") or []]
+
+    @staticmethod
+    def _extracted_from_response(response: dict[str, Any]) -> ExtractedMemory:
+        return ExtractedMemory(
+            name=response["name"],
+            actor_id=response.get("actor_id"),
+            path=response.get("path"),
+            content=response.get("content"),
+            description=response.get("description"),
+            session_id=response.get("session_id"),
+            source_type=response.get("source_type"),
+            create_time=parse_timestamp(response.get("create_time")),
+            update_time=parse_timestamp(response.get("update_time")),
+        )
+
+    def _store_from_response(self, response: dict[str, Any]) -> SessionStore:
+        name = response.get("session_store_name")
+        if name is None:
+            name = response["name"].removeprefix("session-stores/")
+        return SessionStore(
+            name=name,
+            session_store_id=response.get("session_store_id"),
+            creator_user_id=response.get("creator_user_id"),
+            create_time=parse_timestamp(response.get("create_time")),
+            update_time=parse_timestamp(response.get("update_time")),
+            description=response.get("description"),
+            metadata=dict(response.get("metadata", {})),
+            _client=self,
+        )
+
+    def _session_from_response(
+        self,
+        response: dict[str, Any],
+        session_store_name: str,
+    ) -> Session:
+        return Session(
+            store_name=response.get("session_store_name", session_store_name),
+            session_id=response["session_id"],
+            actor_id=response["actor_id"],
+            parent_session_id=response.get("parent_session_id"),
+            root_session_id=response.get("root_session_id"),
+            metadata=dict(response.get("metadata", {})),
+            create_time=parse_timestamp(response.get("create_time")),
+            update_time=parse_timestamp(response.get("update_time")),
+            last_activity_time=parse_timestamp(response.get("last_activity_time")),
+            _client=self,
+        )
+
+    @staticmethod
+    def _item_from_response(response: dict[str, Any]) -> SessionItem:
+        return SessionItem(
+            item_id=response["item_id"],
+            data=response["data"],
+            create_time=parse_timestamp(response.get("create_time")),
+        )

databricks_agentkit/timefmt.py

@@ -0,0 +1,85 @@
+"""Timestamp parsing and humanization for the Agent Bricks CLI.
+
+The 2.0 agents APIs return timestamps as `google.protobuf.Timestamp`, which
+serializes to an RFC 3339 string (`2026-08-15T01:29:00Z`) under the fields
+`create_time`/`update_time`. Older memory-store responses used epoch-millis
+int64 (`created_at`/`updated_at`); `parse_timestamp` accepts either form.
+
+`relative` returns concise phrases such as "13 days ago" and "An hour ago".
+"""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import Optional, Union
+
+TimestampValue = Union[int, float, str, datetime, None]
+
+
+def parse_timestamp(value: TimestampValue) -> Optional[datetime]:
+    """Parse a datetime, an epoch-millis number, a numeric string, or an RFC 3339 string.
+
+    Returns a timezone-aware UTC datetime, or None if the value is empty/unparseable.
+    """
+    if value is None or value == "":
+        return None
+    if isinstance(value, datetime):
+        return value if value.tzinfo else value.replace(tzinfo=timezone.utc)
+    if isinstance(value, (int, float)):
+        return datetime.fromtimestamp(value / 1000.0, tz=timezone.utc)
+    if isinstance(value, str):
+        if value.isdigit():
+            return datetime.fromtimestamp(int(value) / 1000.0, tz=timezone.utc)
+        iso = value.replace("Z", "+00:00")
+        try:
+            dt = datetime.fromisoformat(iso)
+        except ValueError:
+            return None
+        return dt if dt.tzinfo else dt.replace(tzinfo=timezone.utc)
+    return None
+
+
+def relative(value: TimestampValue, *, now: Optional[datetime] = None) -> str:
+    """Humanize a timestamp as an "N units ago" string."""
+    dt = parse_timestamp(value)
+    if dt is None:
+        return "—"
+    now = now or datetime.now(tz=timezone.utc)
+    seconds = (now - dt).total_seconds()
+    if seconds < 0:
+        return "just now"
+
+    minutes = seconds / 60
+    hours = minutes / 60
+    days = hours / 24
+
+    if seconds < 45:
+        return "just now"
+    if minutes < 2:
+        return "A minute ago"
+    if minutes < 60:
+        return f"{round(minutes)} minutes ago"
+    if hours < 2:
+        return "An hour ago"
+    if hours < 24:
+        return f"{round(hours)} hours ago"
+    if days < 2:
+        return "A day ago"
+    if days < 30:
+        return f"{round(days)} days ago"
+    if days < 60:
+        return "A month ago"
+    if days < 365:
+        return f"{round(days / 30)} months ago"
+    if days < 730:
+        return "A year ago"
+    return f"{round(days / 365)} years ago"
+
+
+def absolute(value: TimestampValue) -> str:
+    """Format a timestamp like the mock's detail rail: "May 19, 2026, 05:00 PM"."""
+    dt = parse_timestamp(value)
+    if dt is None:
+        return "—"
+    local = dt.astimezone()
+    return f"{local.strftime('%b')} {local.day}, {local.strftime('%Y, %I:%M %p')}"