community-agent-framework-postgres 0.1.1__py3-none-any.whl

agent_framework_community_postgres/_session_store.py

@@ -0,0 +1,135 @@
+"""``AgentSession`` snapshots in PostgreSQL, one row per opaque session-store id."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any, ClassVar, cast
+
+from agent_framework import AgentSession, SecretString, SessionStore
+from psycopg import sql
+from psycopg.rows import dict_row
+
+from ._client import ClientHandle, PostgresClient
+from ._entra import EntraCredential
+from ._json import encode_jsonb
+from ._retention import EXPIRES_AT, PurgeReport, RetentionPolicy, purge_rows
+from ._store import BaseStore
+
+
+class PostgresSessionStore(SessionStore, BaseStore):
+    """PostgreSQL-backed ``SessionStore``.
+
+    ``set`` stores ``AgentSession.to_dict()`` (custom state goes through
+    ``register_state_type`` codecs) as JSONB; ``get`` rebuilds a fresh session with
+    ``AgentSession.from_dict``. Restore a session with the same agent and provider
+    configuration that created it, as the framework documents. ``SessionStore`` is
+    marked experimental upstream; this implementation follows it.
+    Writes are last-writer-wins, as with the framework's ``SessionStore``; serialize
+    concurrent runs of one session in the host.
+    """
+
+    SNAPSHOT_VERSION: ClassVar[str] = "1"
+
+    def __init__(
+        self,
+        *,
+        application_id: str,
+        connection_string: str | SecretString | None = None,
+        client: PostgresClient | ClientHandle | None = None,
+        credential: EntraCredential | None = None,
+        env_file_path: str | None = None,
+        env_file_encoding: str | None = None,
+        schema: str = "public",
+        table_prefix: str = "af_",
+        retention: RetentionPolicy | None = None,
+        pool_kwargs: Mapping[str, Any] | None = None,
+    ) -> None:
+        SessionStore.__init__(self)
+        BaseStore.__init__(
+            self,
+            application_id=application_id,
+            connection_string=connection_string,
+            client=client,
+            credential=credential,
+            env_file_path=env_file_path,
+            env_file_encoding=env_file_encoding,
+            schema=schema,
+            table_prefix=table_prefix,
+            retention=retention,
+            pool_kwargs=pool_kwargs,
+        )
+
+    async def get(self, session_id: str) -> AgentSession | None:
+        """A fresh copy of the stored session, or ``None`` when absent or purged.
+
+        Raises:
+            ValueError: The row holds a snapshot this release cannot read. Delete it to start over.
+        """
+        SessionStore.validate_session_id(session_id)
+        async with (
+            self._client.connection(transactional=False) as connection,
+            connection.cursor(row_factory=dict_row) as cursor,
+        ):
+            await cursor.execute(
+                sql.SQL(
+                    "SELECT snapshot, snapshot_version, purged_at FROM {sessions}"
+                    " WHERE application_id = %(app)s AND session_id = %(id)s"
+                ).format(sessions=self._names.table("sessions")),
+                {"app": self.application_id, "id": session_id},
+            )
+            row = await cursor.fetchone()
+        if row is None or row["purged_at"] is not None:
+            return None
+        if row["snapshot_version"] != self.SNAPSHOT_VERSION:
+            raise ValueError(
+                f"Unsupported session snapshot version {row['snapshot_version']!r}; expected {self.SNAPSHOT_VERSION!r}."
+            )
+        if not isinstance(row["snapshot"], dict):
+            raise ValueError("The stored session snapshot is not an object.")
+        return AgentSession.from_dict(cast(dict[str, Any], row["snapshot"]))
+
+    async def set(self, session_id: str, session: AgentSession) -> None:
+        """Store the session, replacing any existing snapshot."""
+        SessionStore.validate_session_id(session_id)
+        snapshot = encode_jsonb(session.to_dict(), what="Session state")
+        async with self._client.connection() as connection:
+            await connection.execute(
+                sql.SQL(
+                    "INSERT INTO {sessions} (application_id, session_id, snapshot, snapshot_version, expires_at)"
+                    " VALUES (%(app)s, %(id)s, %(snapshot)s, %(version)s, {expires})"
+                    " ON CONFLICT (application_id, session_id) DO UPDATE SET"
+                    " snapshot = EXCLUDED.snapshot, snapshot_version = EXCLUDED.snapshot_version,"
+                    " updated_at = now(),"
+                    " expires_at = EXCLUDED.expires_at, purged_at = NULL"
+                ).format(sessions=self._names.table("sessions"), expires=EXPIRES_AT),
+                {
+                    "app": self.application_id,
+                    "id": session_id,
+                    "snapshot": snapshot,
+                    "version": self.SNAPSHOT_VERSION,
+                    "ttl": self._ttl(),
+                },
+            )
+
+    async def delete(self, session_id: str) -> None:
+        """Remove the stored session, if any."""
+        SessionStore.validate_session_id(session_id)
+        async with self._client.connection() as connection:
+            await connection.execute(
+                sql.SQL("DELETE FROM {sessions} WHERE application_id = %(app)s AND session_id = %(id)s").format(
+                    sessions=self._names.table("sessions")
+                ),
+                {"app": self.application_id, "id": session_id},
+            )
+
+    async def purge(self) -> PurgeReport:
+        """Purge expired sessions in this store's mode, whatever policy stamped them."""
+        count = await purge_rows(
+            self._client,
+            table=self._names.table("sessions"),
+            where=sql.SQL("application_id = %(app)s"),
+            params={"app": self.application_id},
+            mode=self._retention.mode,
+            tombstone=sql.SQL("snapshot = NULL"),
+        )
+        return PurgeReport({f"{self._names.prefix}sessions": count})

agent_framework_community_postgres/_store.py

@@ -0,0 +1,66 @@
+"""What every store shares: one client, one naming scheme, one retention policy."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from datetime import timedelta
+from typing import Any, Self
+
+from agent_framework import SecretString
+
+from ._client import (
+    ClientHandle,
+    PostgresClient,
+    TableNames,
+    create_client,
+    require_text,
+)
+from ._entra import EntraCredential
+from ._framework import log_untested_once
+from ._retention import RetentionPolicy
+
+
+class BaseStore:
+    def __init__(
+        self,
+        *,
+        application_id: str,
+        connection_string: str | SecretString | None,
+        client: PostgresClient | ClientHandle | None,
+        credential: EntraCredential | None = None,
+        env_file_path: str | None,
+        env_file_encoding: str | None,
+        schema: str,
+        table_prefix: str,
+        retention: RetentionPolicy | None,
+        pool_kwargs: Mapping[str, Any] | None = None,
+    ) -> None:
+        log_untested_once()
+        self.application_id = require_text(application_id, "application_id")
+        self._names = TableNames(schema=schema, prefix=table_prefix)
+        self._client: ClientHandle = create_client(
+            connection_string,
+            client=client,
+            env_file_path=env_file_path,
+            env_file_encoding=env_file_encoding,
+            credential=credential,
+            pool_kwargs=pool_kwargs,
+        )
+        self._retention = retention or RetentionPolicy()
+
+    @property
+    def retention(self) -> RetentionPolicy:
+        return self._retention
+
+    def _ttl(self) -> timedelta | None:
+        return self._retention.ttl
+
+    async def close(self) -> None:
+        """Close the pool this store created; a borrowed client is left alone."""
+        await self._client.close()
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        await self.close()

agent_framework_community_postgres/_thread_snapshot_store.py

@@ -0,0 +1,162 @@
+"""Latest AG-UI Thread Snapshots in PostgreSQL, keyed by scope and thread id."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any
+
+from agent_framework import SecretString
+from psycopg import sql
+from psycopg.rows import dict_row
+from psycopg.types.json import Jsonb
+
+from ._client import ClientHandle, PostgresClient, require_text
+from ._entra import EntraCredential
+from ._json import encode_jsonb
+from ._retention import EXPIRES_AT, PurgeReport, RetentionPolicy, purge_rows
+from ._store import BaseStore
+
+try:
+    from agent_framework_ag_ui import AGUIThreadSnapshot, AGUIThreadSnapshotStore, SnapshotScope
+except ImportError as exc:  # pragma: no cover - exercised only without the extra
+    raise ImportError(
+        "PostgresAGUIThreadSnapshotStore needs agent-framework-ag-ui. "
+        "Install community-agent-framework-postgres[ag-ui]."
+    ) from exc
+
+
+class PostgresAGUIThreadSnapshotStore(BaseStore):
+    """PostgreSQL-backed ``AGUIThreadSnapshotStore``.
+
+    One row per ``(scope, thread_id)`` holding the latest snapshot; ``save`` replaces it
+    (last writer wins, as the in-memory store documents). ``scope`` must be the
+    application's authorization boundary; a thread id alone never reads a row.
+    """
+
+    def __init__(
+        self,
+        *,
+        application_id: str,
+        connection_string: str | SecretString | None = None,
+        client: PostgresClient | ClientHandle | None = None,
+        credential: EntraCredential | None = None,
+        env_file_path: str | None = None,
+        env_file_encoding: str | None = None,
+        schema: str = "public",
+        table_prefix: str = "af_",
+        retention: RetentionPolicy | None = None,
+        pool_kwargs: Mapping[str, Any] | None = None,
+    ) -> None:
+        super().__init__(
+            application_id=application_id,
+            connection_string=connection_string,
+            client=client,
+            credential=credential,
+            env_file_path=env_file_path,
+            env_file_encoding=env_file_encoding,
+            schema=schema,
+            table_prefix=table_prefix,
+            retention=retention,
+            pool_kwargs=pool_kwargs,
+        )
+
+    def _params(self, scope: str, thread_id: str) -> dict[str, Any]:
+        return {
+            "app": self.application_id,
+            "scope": require_text(scope, "scope"),
+            "thread": require_text(thread_id, "thread_id"),
+        }
+
+    _KEY = sql.SQL("application_id = %(app)s AND scope = %(scope)s AND thread_id = %(thread)s")
+
+    @staticmethod
+    def _jsonb(value: Any, what: str) -> Jsonb | None:
+        return encode_jsonb(value, what=what) if value is not None else None
+
+    async def save(self, *, scope: str, thread_id: str, snapshot: AGUIThreadSnapshot) -> None:
+        """Store the latest snapshot for the thread, replacing the previous one."""
+        values = {
+            "messages": encode_jsonb(list(snapshot.messages), what="Thread snapshot messages"),
+            "state": self._jsonb(snapshot.state, "Thread snapshot state"),
+            "interrupt": self._jsonb(snapshot.interrupt, "Thread snapshot interrupt"),
+            "session_state": self._jsonb(snapshot.session_state, "Thread snapshot session state"),
+        }
+        async with self._client.connection() as connection:
+            await connection.execute(
+                sql.SQL(
+                    "INSERT INTO {snapshots} (application_id, scope, thread_id, messages, state, interrupt,"
+                    " session_state, expires_at)"
+                    " VALUES (%(app)s, %(scope)s, %(thread)s, %(messages)s, %(state)s, %(interrupt)s,"
+                    " %(session_state)s, {expires})"
+                    " ON CONFLICT (application_id, scope, thread_id) DO UPDATE SET"
+                    " messages = EXCLUDED.messages, state = EXCLUDED.state, interrupt = EXCLUDED.interrupt,"
+                    " session_state = EXCLUDED.session_state,"
+                    " updated_at = now(), expires_at = EXCLUDED.expires_at, purged_at = NULL"
+                ).format(snapshots=self._names.table("thread_snapshots"), expires=EXPIRES_AT),
+                self._params(scope, thread_id) | values | {"ttl": self._ttl()},
+            )
+
+    async def get(self, *, scope: str, thread_id: str) -> AGUIThreadSnapshot | None:
+        """The latest snapshot, or ``None`` when absent or purged."""
+        async with (
+            self._client.connection(transactional=False) as connection,
+            connection.cursor(row_factory=dict_row) as cursor,
+        ):
+            await cursor.execute(
+                sql.SQL(
+                    "SELECT messages, state, interrupt, session_state, purged_at FROM {snapshots} WHERE {key}"
+                ).format(snapshots=self._names.table("thread_snapshots"), key=self._KEY),
+                self._params(scope, thread_id),
+            )
+            row = await cursor.fetchone()
+        if row is None or row["purged_at"] is not None:
+            return None
+        return AGUIThreadSnapshot(
+            messages=row["messages"], state=row["state"], interrupt=row["interrupt"], session_state=row["session_state"]
+        )
+
+    async def delete(self, *, scope: str, thread_id: str) -> bool:
+        """Remove the thread's snapshot; ``True`` when a live one existed.
+
+        A purged row is removed too, and returns ``False``.
+        """
+        async with self._client.connection() as connection:
+            cursor = await connection.execute(
+                sql.SQL("DELETE FROM {snapshots} WHERE {key} RETURNING purged_at IS NULL").format(
+                    snapshots=self._names.table("thread_snapshots"), key=self._KEY
+                ),
+                self._params(scope, thread_id),
+            )
+            row = await cursor.fetchone()
+        return row is not None and bool(row[0])
+
+    async def clear(self, *, scope: SnapshotScope | None = None) -> None:
+        """Remove one scope's snapshots, or every snapshot of this application when ``scope`` is ``None``.
+
+        The signature is the ``AGUIThreadSnapshotStore`` protocol's. An explicit empty scope raises ``ValueError``.
+        """
+        if scope is None:
+            query = sql.SQL("DELETE FROM {snapshots} WHERE application_id = %(app)s")
+            params: dict[str, Any] = {"app": self.application_id}
+        else:
+            query = sql.SQL("DELETE FROM {snapshots} WHERE application_id = %(app)s AND scope = %(scope)s")
+            params = {"app": self.application_id, "scope": require_text(scope, "scope")}
+        async with self._client.connection() as connection:
+            await connection.execute(query.format(snapshots=self._names.table("thread_snapshots")), params)
+
+    async def purge(self) -> PurgeReport:
+        """Purge expired snapshots in this store's mode, whatever policy stamped them."""
+        count = await purge_rows(
+            self._client,
+            table=self._names.table("thread_snapshots"),
+            where=sql.SQL("application_id = %(app)s"),
+            params={"app": self.application_id},
+            mode=self._retention.mode,
+            tombstone=sql.SQL("messages = NULL, state = NULL, interrupt = NULL, session_state = NULL"),
+        )
+        return PurgeReport({f"{self._names.prefix}thread_snapshots": count})
+
+
+if TYPE_CHECKING:
+    # The assignment is the test; pyright checks protocol compatibility.
+    _PROTOCOL_CHECK: type[AGUIThreadSnapshotStore] = PostgresAGUIThreadSnapshotStore