databricks-agentbricks 0.2.0__py3-none-any.whl

databricks_agentkit/memory_store.py

@@ -0,0 +1,371 @@
+"""Typed managed-memory resources and operations."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import TYPE_CHECKING, Any, List, Optional
+
+from databricks_agentkit._pagination import validate_limit, validate_page_size
+from databricks_agentkit.timefmt import parse_timestamp
+
+if TYPE_CHECKING:
+    from databricks_agentkit._api_client import _AgentBricksApiClient
+
+
+def _resource_id(name: str) -> str:
+    return name.rsplit("/", 1)[-1]
+
+
+@dataclass(frozen=True, kw_only=True)
+class Memory:
+    name: str
+    actor_id: str
+    path: str
+    session_id: Optional[str] = None
+    content: Optional[str] = None
+    description: Optional[str] = None
+    source_type: Optional[str] = None
+    create_time: Optional[datetime] = None
+    update_time: Optional[datetime] = None
+    _store: MemoryStore = field(repr=False, compare=False)
+    _client: MemoryStores = field(repr=False, compare=False)
+
+    @property
+    def id(self) -> str:
+        return _resource_id(self.name)
+
+    def update(
+        self,
+        *,
+        content: Optional[str] = None,
+        description: Optional[str] = None,
+    ) -> Memory:
+        return self._client._update_memory(
+            self._store,
+            memory_id=self.id,
+            content=content,
+            description=description,
+        )
+
+    def delete(self) -> None:
+        self._client._delete_memory(self._store, memory_id=self.id)
+
+
+@dataclass(frozen=True)
+class MemorySearchResult:
+    memory: Memory
+    score: Optional[float] = None
+
+
+@dataclass(frozen=True, kw_only=True)
+class MemoryStore:
+    name: str
+    display_name: str
+    workspace_id: Optional[int] = None
+    storage_backend: Optional[dict[str, Any]] = None
+    owner_user_id: Optional[str] = None
+    create_time: Optional[datetime] = None
+    update_time: Optional[datetime] = None
+    description: Optional[str] = None
+    _client: MemoryStores = field(repr=False, compare=False)
+
+    @property
+    def id(self) -> str:
+        return _resource_id(self.name)
+
+    def update(
+        self,
+        *,
+        display_name: Optional[str] = None,
+        description: Optional[str] = None,
+    ) -> MemoryStore:
+        return self._client._update_store(
+            self,
+            display_name=display_name,
+            description=description,
+        )
+
+    def delete(self) -> None:
+        self._client._delete_store(self)
+
+    def add(
+        self,
+        *,
+        actor_id: str,
+        path: str,
+        content: str,
+        session_id: Optional[str] = None,
+        description: Optional[str] = None,
+        source_type: Optional[str] = None,
+        write_mode: Optional[str] = None,
+    ) -> Memory:
+        return self._client._add_memory(
+            self,
+            actor_id=actor_id,
+            path=path,
+            content=content,
+            session_id=session_id,
+            description=description,
+            source_type=source_type,
+            write_mode=write_mode,
+        )
+
+    def list(
+        self,
+        *,
+        actor_id: str,
+        path_prefix: Optional[str] = None,
+        session_id: Optional[str] = None,
+        page_size: Optional[int] = None,
+        read_mask: Optional[str] = None,
+    ) -> Iterator[Memory]:
+        return self._client._list_memories(
+            self,
+            actor_id=actor_id,
+            path_prefix=path_prefix,
+            session_id=session_id,
+            page_size=page_size,
+            read_mask=read_mask,
+        )
+
+    def get(
+        self,
+        memory_id: str,
+        *,
+        read_mask: Optional[str] = None,
+    ) -> Memory:
+        return self._client._get_memory(self, memory_id=memory_id, read_mask=read_mask)
+
+    def search(
+        self,
+        *,
+        actor_id: str,
+        query: str,
+        limit: Optional[int] = None,
+        path_prefix: Optional[str] = None,
+        session_id: Optional[str] = None,
+        read_mask: Optional[str] = None,
+    ) -> List[MemorySearchResult]:
+        return self._client._search_memories(
+            self,
+            actor_id=actor_id,
+            query=query,
+            limit=limit,
+            path_prefix=path_prefix,
+            session_id=session_id,
+            read_mask=read_mask,
+        )
+
+    def grant_permission(self, principal_id: str, *, permission: str = "WRITE") -> None:
+        """Grant a service principal READ/WRITE access to this memory 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 MemoryStores:
+    def __init__(self, api: _AgentBricksApiClient):
+        self._api = api
+
+    def create(
+        self,
+        display_name: str,
+        *,
+        description: Optional[str] = None,
+    ) -> MemoryStore:
+        response = self._api.create_memory_store(display_name, description)
+        return self._store_from_response(response)
+
+    def list(self, *, page_size: Optional[int] = None) -> Iterator[MemoryStore]:
+        validate_page_size(page_size)
+        page_token = None
+        while True:
+            response = self._api.list_memory_stores(
+                page_size=page_size,
+                page_token=page_token,
+            )
+            for store in response.get("managed_memory_stores", []):
+                yield self._store_from_response(store)
+            page_token = response.get("next_page_token")
+            if not page_token:
+                return
+
+    def get(self, store_id: str) -> MemoryStore:
+        return self._store_from_response(self._api.get_memory_store(store_id))
+
+    def _update_store(
+        self,
+        store: MemoryStore,
+        *,
+        display_name: Optional[str] = None,
+        description: Optional[str] = None,
+    ) -> MemoryStore:
+        if display_name is None and description is None:
+            raise ValueError("at least one of display_name and description is required")
+        response = self._api.update_memory_store(
+            store.id,
+            display_name=display_name,
+            description=description,
+        )
+        return self._store_from_response(response)
+
+    def _delete_store(self, store: MemoryStore) -> None:
+        self._api.delete_memory_store(store.id)
+
+    def _grant_permission(self, store: MemoryStore, principal_id: str, *, permission: str) -> None:
+        self._api.grant_memory_store_permission(store.id, principal_id, permission=permission)
+
+    def _add_memory(
+        self,
+        store: MemoryStore,
+        *,
+        actor_id: str,
+        path: str,
+        content: str,
+        session_id: Optional[str] = None,
+        description: Optional[str] = None,
+        source_type: Optional[str] = None,
+        write_mode: Optional[str] = None,
+    ) -> Memory:
+        response = self._api.create_memory_entry(
+            store.id,
+            actor_id,
+            path,
+            content=content,
+            description=description,
+            session_id=session_id,
+            source_type=source_type,
+            write_mode=write_mode,
+        )
+        return self._memory_from_response(response, store)
+
+    def _list_memories(
+        self,
+        store: MemoryStore,
+        *,
+        actor_id: str,
+        path_prefix: Optional[str] = None,
+        session_id: Optional[str] = None,
+        page_size: Optional[int] = None,
+        read_mask: Optional[str] = None,
+    ) -> Iterator[Memory]:
+        validate_page_size(page_size)
+        page_token = None
+        while True:
+            response = self._api.list_memory_entries(
+                store.id,
+                actor_id,
+                path_prefix=path_prefix,
+                session_id=session_id,
+                page_size=page_size,
+                page_token=page_token,
+                read_mask=read_mask,
+            )
+            for memory in response.get("managed_memory_entries", []):
+                yield self._memory_from_response(memory, store)
+            page_token = response.get("next_page_token")
+            if not page_token:
+                return
+
+    def _get_memory(
+        self,
+        store: MemoryStore,
+        *,
+        memory_id: str,
+        read_mask: Optional[str] = None,
+    ) -> Memory:
+        response = self._api.get_memory_entry(
+            store.id,
+            memory_id,
+            read_mask=read_mask,
+        )
+        return self._memory_from_response(response, store)
+
+    def _search_memories(
+        self,
+        store: MemoryStore,
+        *,
+        actor_id: str,
+        query: str,
+        limit: Optional[int] = None,
+        path_prefix: Optional[str] = None,
+        session_id: Optional[str] = None,
+        read_mask: Optional[str] = None,
+    ) -> List[MemorySearchResult]:
+        effective_limit = 100 if limit is None else limit
+        validate_limit(effective_limit)
+        response = self._api.search_memory_entries(
+            store.id,
+            actor_id,
+            query,
+            page_size=effective_limit,
+            path_prefix=path_prefix,
+            session_id=session_id,
+            read_mask=read_mask,
+        )
+        results = response.get("results")
+        if results is not None:
+            return [
+                MemorySearchResult(
+                    memory=self._memory_from_response(result["managed_memory_entry"], store),
+                    score=result.get("score"),
+                )
+                for result in results
+            ]
+        return [
+            MemorySearchResult(memory=self._memory_from_response(memory, store))
+            for memory in response.get("managed_memory_entries", [])
+        ]
+
+    def _update_memory(
+        self,
+        store: MemoryStore,
+        *,
+        memory_id: str,
+        content: Optional[str] = None,
+        description: Optional[str] = None,
+    ) -> Memory:
+        if content is None and description is None:
+            raise ValueError("at least one of content and description is required")
+        response = self._api.update_memory_entry(
+            store.id,
+            memory_id,
+            content=content,
+            description=description,
+        )
+        return self._memory_from_response(response, store)
+
+    def _delete_memory(self, store: MemoryStore, *, memory_id: str) -> None:
+        self._api.delete_memory_entry(store.id, memory_id)
+
+    def _store_from_response(self, response: dict[str, Any]) -> MemoryStore:
+        return MemoryStore(
+            name=response["name"],
+            display_name=response["display_name"],
+            workspace_id=response.get("workspace_id"),
+            storage_backend=response.get("storage_backend"),
+            owner_user_id=response.get("owner_user_id"),
+            create_time=parse_timestamp(response.get("created_at") or response.get("create_time")),
+            update_time=parse_timestamp(response.get("updated_at") or response.get("update_time")),
+            description=response.get("description"),
+            _client=self,
+        )
+
+    def _memory_from_response(self, response: dict[str, Any], store: MemoryStore) -> Memory:
+        return Memory(
+            name=response["name"],
+            actor_id=response["actor_id"],
+            session_id=response.get("session_id"),
+            path=response["path"],
+            content=response.get("content"),
+            description=response.get("description"),
+            source_type=response.get("source_type"),
+            create_time=parse_timestamp(response.get("create_time")),
+            update_time=parse_timestamp(response.get("update_time")),
+            _store=store,
+            _client=self,
+        )

databricks_agentkit/models.py

@@ -0,0 +1,330 @@
+"""Typed, forward-compatible views over 2.0 agents JSON responses.
+
+Every model subclasses `dict` and holds the complete raw response, so mapping
+access (`store["name"]`), `json.dumps`, and any new/unknown server fields keep
+working. Subclasses add typed `@property` accessors as a convenience layer;
+they never mutate the stored data, so nothing is lost as the preview API grows.
+
+Accessors are intentionally named to avoid shadowing `dict` methods (e.g. list
+models expose `.stores` / `.session_items`, never `.items`).
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+
+class _Model(dict):
+    """A dict-backed API response with typed accessors over the raw JSON."""
+
+    def to_dict(self) -> dict:
+        """A plain-dict shallow copy of the raw response."""
+        return dict(self)
+
+
+def _wrap(value: Any, cls: type) -> Any:
+    return cls(value) if isinstance(value, dict) else None
+
+
+def _wrap_list(value: Any, cls: type) -> list:
+    return [cls(v) for v in value if isinstance(v, dict)] if isinstance(value, list) else []
+
+
+class StorageBackend(_Model):
+    @property
+    def backend_type(self) -> Optional[str]:
+        return self.get("backend_type")
+
+    @property
+    def backend_id(self) -> Optional[str]:
+        return self.get("backend_id")
+
+
+class RuntimeStoreLakebaseBackend(_Model):
+    @property
+    def project_id(self) -> Optional[str]:
+        return self.get("project_id")
+
+    @property
+    def branch(self) -> Optional[str]:
+        return self.get("branch")
+
+    @property
+    def database_id(self) -> Optional[str]:
+        return self.get("database_id")
+
+
+class RuntimeStoreAppOwner(_Model):
+    @property
+    def name(self) -> Optional[str]:
+        return self.get("name")
+
+    @property
+    def service_principal_id(self) -> Optional[str]:
+        return self.get("service_principal_id")
+
+
+class RuntimeStoreOwner(_Model):
+    @property
+    def app(self) -> Optional[RuntimeStoreAppOwner]:
+        return _wrap(self.get("app"), RuntimeStoreAppOwner)
+
+
+class RuntimeStoreStorageBackend(_Model):
+    @property
+    def lakebase(self) -> Optional[RuntimeStoreLakebaseBackend]:
+        return _wrap(self.get("lakebase"), RuntimeStoreLakebaseBackend)
+
+
+class RuntimeStore(_Model):
+    @property
+    def name(self) -> Optional[str]:
+        return self.get("name")
+
+    @property
+    def owner(self) -> Optional[RuntimeStoreOwner]:
+        return _wrap(self.get("owner"), RuntimeStoreOwner)
+
+    @property
+    def storage_backend(self) -> Optional[RuntimeStoreStorageBackend]:
+        return _wrap(self.get("storage_backend"), RuntimeStoreStorageBackend)
+
+
+class MemoryStore(_Model):
+    @property
+    def name(self) -> Optional[str]:
+        return self.get("name")
+
+    @property
+    def display_name(self) -> Optional[str]:
+        return self.get("display_name")
+
+    @property
+    def workspace_id(self) -> Optional[str]:
+        return self.get("workspace_id")
+
+    @property
+    def owner_user_id(self) -> Optional[str]:
+        return self.get("owner_user_id")
+
+    @property
+    def description(self) -> Optional[str]:
+        return self.get("description")
+
+    @property
+    def storage_backend(self) -> Optional[StorageBackend]:
+        return _wrap(self.get("storage_backend"), StorageBackend)
+
+    @property
+    def create_time(self) -> Optional[str]:
+        return self.get("create_time")
+
+    @property
+    def update_time(self) -> Optional[str]:
+        return self.get("update_time")
+
+
+class MemoryEntry(_Model):
+    @property
+    def name(self) -> Optional[str]:
+        return self.get("name")
+
+    @property
+    def actor_id(self) -> Optional[str]:
+        return self.get("actor_id")
+
+    @property
+    def path(self) -> Optional[str]:
+        return self.get("path")
+
+    @property
+    def content(self) -> Optional[str]:
+        return self.get("content")
+
+    @property
+    def description(self) -> Optional[str]:
+        return self.get("description")
+
+    @property
+    def session_id(self) -> Optional[str]:
+        return self.get("session_id")
+
+    @property
+    def source_type(self) -> Optional[str]:
+        return self.get("source_type")
+
+    @property
+    def create_time(self) -> Optional[str]:
+        return self.get("create_time")
+
+    @property
+    def update_time(self) -> Optional[str]:
+        return self.get("update_time")
+
+
+class MemoryStoreList(_Model):
+    @property
+    def stores(self) -> list[MemoryStore]:
+        return _wrap_list(self.get("managed_memory_stores"), MemoryStore)
+
+    @property
+    def next_page_token(self) -> Optional[str]:
+        return self.get("next_page_token")
+
+
+class MemoryEntryList(_Model):
+    @property
+    def entries(self) -> list[MemoryEntry]:
+        return _wrap_list(self.get("managed_memory_entries"), MemoryEntry)
+
+    @property
+    def next_page_token(self) -> Optional[str]:
+        return self.get("next_page_token")
+
+
+class ExtractMemoriesResponse(_Model):
+    @property
+    def name(self) -> Optional[str]:
+        return self.get("name")
+
+    @property
+    def entries(self) -> list[MemoryEntry]:
+        return _wrap_list(self.get("entries"), MemoryEntry)
+
+
+class MemorySearchHit(_Model):
+    @property
+    def entry(self) -> Optional[MemoryEntry]:
+        return _wrap(self.get("managed_memory_entry"), MemoryEntry)
+
+    @property
+    def score(self) -> Optional[float]:
+        return self.get("score")
+
+
+class MemorySearchResult(_Model):
+    @property
+    def entries(self) -> list[MemoryEntry]:
+        return _wrap_list(self.get("managed_memory_entries"), MemoryEntry)
+
+    @property
+    def results(self) -> list[MemorySearchHit]:
+        return _wrap_list(self.get("results"), MemorySearchHit)
+
+    @property
+    def next_page_token(self) -> Optional[str]:
+        return self.get("next_page_token")
+
+
+class SessionStore(_Model):
+    @property
+    def session_store_name(self) -> Optional[str]:
+        return self.get("session_store_name")
+
+    @property
+    def session_store_id(self) -> Optional[str]:
+        return self.get("session_store_id")
+
+    @property
+    def creator_user_id(self) -> Optional[str]:
+        return self.get("creator_user_id")
+
+    @property
+    def description(self) -> Optional[str]:
+        return self.get("description")
+
+    @property
+    def create_time(self) -> Optional[str]:
+        return self.get("create_time")
+
+    @property
+    def update_time(self) -> Optional[str]:
+        return self.get("update_time")
+
+
+class SessionStoreList(_Model):
+    @property
+    def stores(self) -> list[SessionStore]:
+        return _wrap_list(self.get("session_stores"), SessionStore)
+
+    @property
+    def next_page_token(self) -> Optional[str]:
+        return self.get("next_page_token")
+
+
+class Session(_Model):
+    @property
+    def session_store_name(self) -> Optional[str]:
+        return self.get("session_store_name")
+
+    @property
+    def session_id(self) -> Optional[str]:
+        return self.get("session_id")
+
+    @property
+    def actor_id(self) -> Optional[str]:
+        return self.get("actor_id")
+
+    @property
+    def root_session_id(self) -> Optional[str]:
+        return self.get("root_session_id")
+
+    @property
+    def parent_session_id(self) -> Optional[str]:
+        return self.get("parent_session_id")
+
+    @property
+    def metadata(self) -> Optional[dict]:
+        return self.get("metadata")
+
+    @property
+    def create_time(self) -> Optional[str]:
+        return self.get("create_time")
+
+    @property
+    def update_time(self) -> Optional[str]:
+        return self.get("update_time")
+
+    @property
+    def last_activity_time(self) -> Optional[str]:
+        return self.get("last_activity_time")
+
+
+class SessionList(_Model):
+    @property
+    def sessions(self) -> list[Session]:
+        return _wrap_list(self.get("sessions"), Session)
+
+    @property
+    def next_page_token(self) -> Optional[str]:
+        return self.get("next_page_token")
+
+
+class SessionItem(_Model):
+    @property
+    def item_id(self) -> Optional[str]:
+        return self.get("item_id")
+
+    @property
+    def data(self) -> Any:
+        return self.get("data")
+
+    @property
+    def create_time(self) -> Optional[str]:
+        return self.get("create_time")
+
+
+class SessionItemList(_Model):
+    @property
+    def session_items(self) -> list[SessionItem]:
+        return _wrap_list(self.get("session_items"), SessionItem)
+
+    @property
+    def next_page_token(self) -> Optional[str]:
+        return self.get("next_page_token")
+
+
+class PoppedSessionItem(_Model):
+    @property
+    def item(self) -> Optional[SessionItem]:
+        return _wrap(self.get("item"), SessionItem)