orchid-storage-postgres 1.0.0__py3-none-any.whl

orchid_storage_postgres/__init__.py

@@ -0,0 +1,48 @@
+"""PostgreSQL storage plugin for the Orchid AI framework.
+
+Provides ``OrchidPostgresChatStorage``, a PostgreSQL visibility
+fragment, and a PostgreSQL checkpointer.  Auto-registers via
+``importlib.metadata`` entry points.
+"""
+
+from __future__ import annotations
+
+import logging
+
+__version__ = "1.0.0"
+
+from .chat_storage import OrchidPostgresChatStorage
+from .visibility import _build_postgres_filter
+
+__all__ = ["OrchidPostgresChatStorage"]
+
+logger = logging.getLogger(__name__)
+
+
+async def _build_postgres_checkpointer(dsn: str):
+    """Build an async PostgreSQL checkpointer from a DSN."""
+    from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+
+    checkpointer = AsyncPostgresSaver.from_conn_string(dsn)
+    await checkpointer.setup()
+    logger.info("[orchid-storage-postgres] Checkpointer ready")
+    return checkpointer
+
+
+def _register() -> None:
+    """Entry-point callable — registers the postgres visibility fragment and checkpointer."""
+    try:
+        from orchid_ai.events.visibility import register_visibility_fragment
+
+        register_visibility_fragment("postgres", _build_postgres_filter)
+        logger.debug("[orchid-storage-postgres] Registered visibility fragment")
+    except ImportError:
+        logger.debug("[orchid-storage-postgres] Skipping visibility fragment (not in this orchid-ai version)")
+
+    try:
+        from orchid_ai.checkpointing.factory import register_checkpointer
+
+        register_checkpointer("postgres", _build_postgres_checkpointer)
+        logger.debug("[orchid-storage-postgres] Registered checkpointer")
+    except ImportError:
+        logger.debug("[orchid-storage-postgres] Skipping checkpointer (not in this orchid-ai version)")

orchid_storage_postgres/chat_storage.py

@@ -0,0 +1,249 @@
+"""
+PostgreSQL chat storage — production-grade :class:`OrchidChatStorage` backend.
+
+Backed by ``asyncpg`` with connection pooling.  Implements every
+:class:`OrchidChatStorage` method using ``$1..$N`` placeholders.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import uuid
+from datetime import datetime
+from typing import Any
+
+from orchid_ai.persistence.base import OrchidChatStorage
+from orchid_ai.persistence.models import OrchidChatMessage, OrchidChatSession, utcnow
+
+from .migrations import PostgresMigrationRunner
+
+logger = logging.getLogger(__name__)
+
+
+class OrchidPostgresChatStorage(OrchidChatStorage):
+    """Async PostgreSQL storage for chat sessions and messages.
+
+    Constructor accepts the DSN via ``dsn`` and an optional
+    ``extra_migrations_package`` (dotted import path) so integrators
+    can append their own migrations after the framework's.
+    """
+
+    def __init__(self, *, dsn: str, extra_migrations_package: str | None = None):
+        self._dsn = dsn
+        self._pool: Any = None
+        self._migrator = PostgresMigrationRunner(
+            extra_migrations_package=extra_migrations_package,
+        )
+
+    # ── Lifecycle ────────────────────────────────────────────
+
+    async def init_db(self) -> None:
+        import asyncpg
+
+        self._pool = await asyncpg.create_pool(self._dsn, min_size=1, max_size=10)
+        async with self._pool.acquire() as conn:
+            await self._migrator.run_up(conn)
+        logger.info("[OrchidChatStorage:postgres] Initialised")
+
+    async def close(self) -> None:
+        if self._pool:
+            await self._pool.close()
+            self._pool = None
+
+    async def _conn(self):
+        if self._pool is None:
+            raise RuntimeError("OrchidPostgresChatStorage: init_db() not called")
+        return await self._pool.acquire()
+
+    # ── Sessions ─────────────────────────────────────────────
+
+    async def create_chat(
+        self,
+        tenant_id: str,
+        user_id: str,
+        title: str = "",
+    ) -> OrchidChatSession:
+        now = utcnow()
+        chat = OrchidChatSession(
+            id=str(uuid.uuid4()),
+            tenant_id=tenant_id,
+            user_id=user_id,
+            title=title or "New chat",
+            created_at=now,
+            updated_at=now,
+        )
+        async with self._pool.acquire() as conn:
+            await conn.execute(
+                "INSERT INTO chat_sessions (id, tenant_id, user_id, title, created_at, updated_at) "
+                "VALUES ($1, $2, $3, $4, $5, $6)",
+                chat.id,
+                chat.tenant_id,
+                chat.user_id,
+                chat.title,
+                now,
+                now,
+            )
+        return chat
+
+    async def list_chats(
+        self,
+        tenant_id: str,
+        user_id: str,
+    ) -> list[OrchidChatSession]:
+        async with self._pool.acquire() as conn:
+            rows = await conn.fetch(
+                "SELECT * FROM chat_sessions WHERE tenant_id = $1 AND user_id = $2 ORDER BY updated_at DESC",
+                tenant_id,
+                user_id,
+            )
+        return [_row_to_session(r) for r in rows]
+
+    async def get_chat(self, chat_id: str) -> OrchidChatSession | None:
+        async with self._pool.acquire() as conn:
+            row = await conn.fetchrow(
+                "SELECT * FROM chat_sessions WHERE id = $1",
+                chat_id,
+            )
+        return _row_to_session(row) if row else None
+
+    async def delete_chat(self, chat_id: str) -> None:
+        async with self._pool.acquire() as conn:
+            await conn.execute("DELETE FROM chat_sessions WHERE id = $1", chat_id)
+
+    async def update_title(self, chat_id: str, title: str) -> None:
+        async with self._pool.acquire() as conn:
+            await conn.execute(
+                "UPDATE chat_sessions SET title = $1, updated_at = $2 WHERE id = $3",
+                title,
+                utcnow(),
+                chat_id,
+            )
+
+    async def mark_shared(self, chat_id: str) -> None:
+        async with self._pool.acquire() as conn:
+            await conn.execute(
+                "UPDATE chat_sessions SET is_shared = TRUE, updated_at = $1 WHERE id = $2",
+                utcnow(),
+                chat_id,
+            )
+
+    # ── Messages ─────────────────────────────────────────────
+
+    async def add_message(
+        self,
+        chat_id: str,
+        role: str,
+        content: str,
+        agents_used: list[str] | None = None,
+        metadata: dict | None = None,
+    ) -> OrchidChatMessage:
+        now = utcnow()
+        msg = OrchidChatMessage(
+            id=str(uuid.uuid4()),
+            chat_id=chat_id,
+            role=role,
+            content=content,
+            agents_used=agents_used or [],
+            created_at=now,
+            metadata=metadata or {},
+        )
+        async with self._pool.acquire() as conn:
+            await conn.execute(
+                "INSERT INTO chat_messages (id, chat_id, role, content, agents_used, created_at, metadata) "
+                "VALUES ($1, $2, $3, $4, $5, $6, $7)",
+                msg.id,
+                msg.chat_id,
+                msg.role,
+                msg.content,
+                json.dumps(msg.agents_used),
+                now,
+                json.dumps(msg.metadata),
+            )
+            await conn.execute(
+                "UPDATE chat_sessions SET updated_at = $1 WHERE id = $2",
+                now,
+                chat_id,
+            )
+        return msg
+
+    async def get_messages(
+        self,
+        chat_id: str,
+        limit: int = 50,
+        offset: int = 0,
+    ) -> list[OrchidChatMessage]:
+        async with self._pool.acquire() as conn:
+            rows = await conn.fetch(
+                "SELECT * FROM chat_messages WHERE chat_id = $1 ORDER BY created_at ASC LIMIT $2 OFFSET $3",
+                chat_id,
+                limit,
+                offset,
+            )
+        return [_row_to_message(r) for r in rows]
+
+    # ── Conversation summaries ───────────────────────────────
+
+    async def get_conversation_summary(self, chat_id: str) -> str | None:
+        async with self._pool.acquire() as conn:
+            row = await conn.fetchrow(
+                "SELECT summary_text FROM conversation_summaries WHERE chat_id = $1",
+                chat_id,
+            )
+        return row["summary_text"] if row else None
+
+    async def save_conversation_summary(self, chat_id: str, summary: str, turn_number: int) -> None:
+        async with self._pool.acquire() as conn:
+            await conn.execute(
+                "INSERT INTO conversation_summaries (chat_id, summary_text, turn_number, updated_at) "
+                "VALUES ($1, $2, $3, $4) "
+                "ON CONFLICT (chat_id) DO UPDATE SET summary_text = $2, turn_number = $3, updated_at = $4",
+                chat_id,
+                summary,
+                turn_number,
+                utcnow(),
+            )
+
+
+# ── Row mappers ──────────────────────────────────────────────
+
+
+def _parse_dt(val: Any) -> datetime:
+    if isinstance(val, datetime):
+        return val
+    if isinstance(val, str):
+        try:
+            return datetime.fromisoformat(val)
+        except (ValueError, TypeError):
+            return utcnow()
+    return utcnow()
+
+
+def _row_to_session(row: Any) -> OrchidChatSession:
+    return OrchidChatSession(
+        id=row["id"],
+        tenant_id=row["tenant_id"],
+        user_id=row["user_id"],
+        title=row["title"],
+        created_at=_parse_dt(row["created_at"]),
+        updated_at=_parse_dt(row["updated_at"]),
+        is_shared=bool(row["is_shared"]),
+    )
+
+
+def _row_to_message(row: Any) -> OrchidChatMessage:
+    agents_used = row["agents_used"]
+    if isinstance(agents_used, str):
+        agents_used = json.loads(agents_used)
+    meta = row["metadata"]
+    if isinstance(meta, str):
+        meta = json.loads(meta)
+    return OrchidChatMessage(
+        id=row["id"],
+        chat_id=row["chat_id"],
+        role=row["role"],
+        content=row["content"],
+        agents_used=agents_used or [],
+        created_at=_parse_dt(row["created_at"]),
+        metadata=meta or {},
+    )

orchid_storage_postgres/migrations/__init__.py

@@ -0,0 +1,44 @@
+"""
+PostgreSQL migration runner.
+
+The runner discovers migrations in ``orchid_storage_postgres.migrations``
+and tracks applied versions in a ``_migrations`` table.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from orchid_ai.persistence.migrations.runner import OrchidMigrationRunner
+
+MIGRATIONS_PACKAGE = "orchid_storage_postgres.migrations"
+
+
+class PostgresMigrationRunner(OrchidMigrationRunner):
+    """PostgreSQL-specific migration tracking."""
+
+    dialect = "postgres"
+    migrations_package = MIGRATIONS_PACKAGE
+
+    async def ensure_migrations_table(self, conn: Any) -> None:
+        await conn.execute("""
+            CREATE TABLE IF NOT EXISTS _migrations (
+                version TEXT PRIMARY KEY,
+                description TEXT NOT NULL DEFAULT '',
+                applied_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
+            )
+        """)
+
+    async def get_applied_versions(self, conn: Any) -> set[str]:
+        rows = await conn.fetch("SELECT version FROM _migrations")
+        return {r["version"] for r in rows}
+
+    async def record_version(self, conn: Any, version: str, description: str) -> None:
+        await conn.execute(
+            "INSERT INTO _migrations (version, description) VALUES ($1, $2)",
+            version,
+            description,
+        )
+
+    async def remove_version(self, conn: Any, version: str) -> None:
+        await conn.execute("DELETE FROM _migrations WHERE version = $1", version)

orchid_storage_postgres/migrations/v001_initial_schema.py

@@ -0,0 +1,48 @@
+"""
+Migration v001 — PostgreSQL initial schema.
+
+Creates every framework-owned table in a single pass.
+"""
+
+from __future__ import annotations
+
+import logging
+
+from orchid_ai.persistence.migrations._schema_ddl import PG_UP
+
+logger = logging.getLogger(__name__)
+
+VERSION = "001"
+DESCRIPTION = "PostgreSQL initial schema (chat, MCP outbound, MCP inbound gateway, events)"
+
+_PG_DOWN = [
+    "DROP TABLE IF EXISTS signal_sources CASCADE",
+    "DROP TABLE IF EXISTS job_runs CASCADE",
+    "DROP TABLE IF EXISTS schedules CASCADE",
+    "DROP TABLE IF EXISTS triggers CASCADE",
+    "DROP TABLE IF EXISTS signal_queue_dead_letter CASCADE",
+    "DROP TABLE IF EXISTS signal_queue CASCADE",
+    "DROP TABLE IF EXISTS signals CASCADE",
+    "DROP TABLE IF EXISTS mcp_gateway_tokens CASCADE",
+    "DROP TABLE IF EXISTS mcp_gateway_auth_codes CASCADE",
+    "DROP TABLE IF EXISTS mcp_gateway_clients CASCADE",
+    "DROP TABLE IF EXISTS mcp_client_registrations CASCADE",
+    "DROP TABLE IF EXISTS mcp_oauth_tokens CASCADE",
+    "DROP TABLE IF EXISTS agent_configs CASCADE",
+    "DROP TABLE IF EXISTS conversation_summaries CASCADE",
+    "DROP TABLE IF EXISTS chat_messages CASCADE",
+    "DROP TABLE IF EXISTS chat_sessions CASCADE",
+]
+
+
+async def up(conn, *, dialect: str = "postgres") -> None:
+    """Apply the PostgreSQL initial schema."""
+    for sql in PG_UP:
+        await conn.execute(sql)
+    logger.info("[orchid-storage-postgres] Migration v001 applied (%d statements)", len(PG_UP))
+
+
+async def down(conn, *, dialect: str = "postgres") -> None:
+    """Roll back the PostgreSQL initial schema."""
+    for sql in _PG_DOWN:
+        await conn.execute(sql)

orchid_storage_postgres/visibility.py

@@ -0,0 +1,35 @@
+"""
+PostgreSQL run-visibility filter fragment.
+
+Registers a ``build_run_filter_clause`` implementation for the ``postgres``
+dialect that uses ``$1..$N`` positional parameters (asyncpg convention).
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from orchid_ai.events.visibility import _Filter  # noqa: PLC2701
+
+
+def _build_postgres_filter(auth: Any) -> _Filter:
+    """Return a ``WHERE`` fragment + positional bind params for PostgreSQL."""
+    tenant_key = getattr(auth, "tenant_key", "default")
+    user_id = getattr(auth, "user_id", "")
+    roles = getattr(auth, "roles", frozenset())
+
+    if "admin" in roles:
+        return _Filter(
+            where="tenant_key = $1",
+            params={"tenant_key": tenant_key},
+        )
+    return _Filter(
+        where=(
+            "tenant_key = $1 AND ("
+            "visibility = 'tenant' "
+            "OR (visibility IN ('actor', 'addressed') "
+            "    AND visibility_user_id = $2)"
+            ")"
+        ),
+        params={"tenant_key": tenant_key, "user_id": user_id},
+    )

orchid_storage_postgres-1.0.0.dist-info/METADATA

@@ -0,0 +1,74 @@
+Metadata-Version: 2.4
+Name: orchid-storage-postgres
+Version: 1.0.0
+Summary: PostgreSQL storage backend plugin for the Orchid AI framework
+Project-URL: Homepage, https://github.com/gadz82/orchid-storage-postgres
+Project-URL: Issues, https://github.com/gadz82/orchid-storage-postgres/issues
+License: MIT
+License-File: LICENSE
+Keywords: agents,ai,asyncpg,postgresql,storage
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Requires-Dist: asyncpg>=0.29.0
+Requires-Dist: langgraph-checkpoint-postgres>=2.0.0
+Requires-Dist: orchid-ai>=1.8.2
+Provides-Extra: dev
+Requires-Dist: pytest-asyncio>=0.25.0; extra == 'dev'
+Requires-Dist: pytest-cov>=6.0; extra == 'dev'
+Requires-Dist: pytest>=8.0; extra == 'dev'
+Requires-Dist: ruff>=0.9.0; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# orchid-storage-postgres
+
+PostgreSQL storage backend plugin for the [Orchid AI](https://github.com/gadz82/orchid) framework.
+
+## What it provides
+
+- `OrchidPostgresChatStorage` — implements `OrchidChatStorage` backed by PostgreSQL (asyncpg)
+- PostgreSQL visibility fragment for `build_run_filter_clause`
+- PostgreSQL schema migration (v001)
+
+## Installation
+
+```bash
+pip install orchid-storage-postgres
+```
+
+## Usage
+
+Reference in your `orchid.yml`:
+
+```yaml
+storage:
+  class: orchid_storage_postgres.chat_storage.OrchidPostgresChatStorage
+  dsn: postgresql://user:pass@localhost:5432/orchid
+```
+
+Or build it programmatically:
+
+```python
+from orchid_storage_postgres import OrchidPostgresChatStorage
+
+storage = OrchidPostgresChatStorage(dsn="postgresql://user:pass@localhost:5432/orchid")
+await storage.init_db()
+```
+
+## Development
+
+```bash
+cd orchid-storage-postgres
+pip install -e ".[dev]"
+pytest tests/ -x
+ruff check orchid_storage_postgres/
+```
+
+## License
+
+MIT

orchid_storage_postgres-1.0.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+orchid_storage_postgres/__init__.py,sha256=PA1dsUPRdVPqiGOuyTX9H_hbaKgLEQUF7qZyKmXdxr8,1699
+orchid_storage_postgres/chat_storage.py,sha256=egOTQIP8EKn57WgNVqxi-DBYkDC9hy62NrvZzAuV5T4,8600
+orchid_storage_postgres/visibility.py,sha256=pUAXFOQD7U6yvB39aiITj8CfVzSx-zxcZZQfzKx1pys,1057
+orchid_storage_postgres/migrations/__init__.py,sha256=vApfq2cIgx41bw7O7wK7hO5ZsH4p6lad2PAc9IAe-0U,1451
+orchid_storage_postgres/migrations/v001_initial_schema.py,sha256=P2g18PASFTtmSjcJcHRc6r2NaGF2cxB4Nop-RVWbxmg,1649
+orchid_storage_postgres-1.0.0.dist-info/METADATA,sha256=IcnFr17TZKEu6cULXgLORBqFHo0spEWL4gndKgIl29Q,2090
+orchid_storage_postgres-1.0.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+orchid_storage_postgres-1.0.0.dist-info/entry_points.txt,sha256=bIUMwNTiu6N0pbG3oKQsKOukdyLpYPrQQx2re3ChlLM,144
+orchid_storage_postgres-1.0.0.dist-info/licenses/LICENSE,sha256=1xiBMms-kdU_t3A21unskfGGda1M2uemoNTaRGgtV_s,1077
+orchid_storage_postgres-1.0.0.dist-info/RECORD,,

orchid_storage_postgres-1.0.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.29.0
+Root-Is-Purelib: true
+Tag: py3-none-any

orchid_storage_postgres-1.0.0.dist-info/entry_points.txt

@@ -0,0 +1,5 @@
+[orchid.checkpointers]
+postgres = orchid_storage_postgres:_register
+
+[orchid.visibility_fragments]
+postgres = orchid_storage_postgres:_register

orchid_storage_postgres-1.0.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Francesco Marchesini
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.