Initial commit: add .gitignore and README

fusionagi/memory/postgres_backend.py

@@ -0,0 +1,231 @@
+"""Optional Postgres backend for memory. Requires: pip install fusionagi[memory]."""
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from fusionagi._logger import logger
+
+
+class MemoryBackend(ABC):
+    """Abstract backend for persistent memory storage."""
+
+    @abstractmethod
+    def store(
+        self,
+        id: str,
+        tenant_id: str,
+        user_id: str,
+        session_id: str,
+        type: str,
+        content: dict[str, Any],
+        metadata: dict[str, Any] | None = None,
+        retention_policy: str = "session",
+    ) -> None:
+        """Store a memory item."""
+        ...
+
+    @abstractmethod
+    def get(self, id: str) -> dict[str, Any] | None:
+        """Get a memory item by id."""
+        ...
+
+    @abstractmethod
+    def query(
+        self,
+        tenant_id: str,
+        user_id: str | None = None,
+        session_id: str | None = None,
+        type: str | None = None,
+        limit: int = 100,
+    ) -> list[dict[str, Any]]:
+        """Query memory items."""
+        ...
+
+
+class InMemoryBackend(MemoryBackend):
+    """In-memory implementation for development."""
+
+    def __init__(self) -> None:
+        self._store: dict[str, dict[str, Any]] = {}
+
+    def store(
+        self,
+        id: str,
+        tenant_id: str,
+        user_id: str,
+        session_id: str,
+        type: str,
+        content: dict[str, Any],
+        metadata: dict[str, Any] | None = None,
+        retention_policy: str = "session",
+    ) -> None:
+        self._store[id] = {
+            "id": id,
+            "tenant_id": tenant_id,
+            "user_id": user_id,
+            "session_id": session_id,
+            "type": type,
+            "content": content,
+            "metadata": metadata or {},
+            "retention_policy": retention_policy,
+        }
+
+    def get(self, id: str) -> dict[str, Any] | None:
+        return self._store.get(id)
+
+    def query(
+        self,
+        tenant_id: str,
+        user_id: str | None = None,
+        session_id: str | None = None,
+        type: str | None = None,
+        limit: int = 100,
+    ) -> list[dict[str, Any]]:
+        out = []
+        for v in self._store.values():
+            if v["tenant_id"] != tenant_id:
+                continue
+            if user_id and v["user_id"] != user_id:
+                continue
+            if session_id and v["session_id"] != session_id:
+                continue
+            if type and v["type"] != type:
+                continue
+            out.append(v)
+            if len(out) >= limit:
+                break
+        return out
+
+
+def create_postgres_backend(connection_string: str) -> MemoryBackend | None:
+    """Create Postgres-backed MemoryBackend when psycopg is available."""
+    try:
+        import psycopg
+    except ImportError:
+        logger.debug("psycopg not installed; use pip install fusionagi[memory]")
+        return None
+
+    return PostgresMemoryBackend(connection_string)
+
+
+class PostgresMemoryBackend(MemoryBackend):
+    """Postgres-backed memory storage."""
+
+    def __init__(self, connection_string: str) -> None:
+        self._conn_str = connection_string
+        self._init_schema()
+
+    def _init_schema(self) -> None:
+        import psycopg
+
+        with psycopg.connect(self._conn_str) as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    """
+                    CREATE TABLE IF NOT EXISTS memory_items (
+                    id TEXT PRIMARY KEY,
+                    tenant_id TEXT NOT NULL,
+                    user_id TEXT NOT NULL,
+                    session_id TEXT NOT NULL,
+                    type TEXT NOT NULL,
+                    content JSONB NOT NULL,
+                    metadata JSONB DEFAULT '{}',
+                    retention_policy TEXT DEFAULT 'session',
+                    created_at TIMESTAMPTZ DEFAULT NOW()
+                )
+                """
+            )
+            conn.commit()
+
+    def store(
+        self,
+        id: str,
+        tenant_id: str,
+        user_id: str,
+        session_id: str,
+        type: str,
+        content: dict[str, Any],
+        metadata: dict[str, Any] | None = None,
+        retention_policy: str = "session",
+    ) -> None:
+        import json
+        import psycopg
+
+        with psycopg.connect(self._conn_str) as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    """
+                    INSERT INTO memory_items (id, tenant_id, user_id, session_id, type, content, metadata, retention_policy)
+                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
+                    ON CONFLICT (id) DO UPDATE SET content = EXCLUDED.content, metadata = EXCLUDED.metadata
+                    """,
+                    (id, tenant_id, user_id, session_id, type, json.dumps(content), json.dumps(metadata or {}), retention_policy),
+                )
+            conn.commit()
+
+    def get(self, id: str) -> dict[str, Any] | None:
+        import json
+        import psycopg
+
+        with psycopg.connect(self._conn_str) as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    "SELECT id, tenant_id, user_id, session_id, type, content, metadata, retention_policy FROM memory_items WHERE id = %s",
+                    (id,),
+                )
+                row = cur.fetchone()
+        if not row:
+            return None
+        return {
+            "id": row[0],
+            "tenant_id": row[1],
+            "user_id": row[2],
+            "session_id": row[3],
+            "type": row[4],
+            "content": json.loads(row[5]) if row[5] else {},
+            "metadata": json.loads(row[6]) if row[6] else {},
+            "retention_policy": row[7],
+        }
+
+    def query(
+        self,
+        tenant_id: str,
+        user_id: str | None = None,
+        session_id: str | None = None,
+        type: str | None = None,
+        limit: int = 100,
+    ) -> list[dict[str, Any]]:
+        import json
+        import psycopg
+
+        q = "SELECT id, tenant_id, user_id, session_id, type, content, metadata, retention_policy FROM memory_items WHERE tenant_id = %s"
+        params: list[Any] = [tenant_id]
+        if user_id:
+            q += " AND user_id = %s"
+            params.append(user_id)
+        if session_id:
+            q += " AND session_id = %s"
+            params.append(session_id)
+        if type:
+            q += " AND type = %s"
+            params.append(type)
+        q += " ORDER BY created_at DESC LIMIT %s"
+        params.append(limit)
+
+        with psycopg.connect(self._conn_str) as conn:
+            with conn.cursor() as cur:
+                cur.execute(q, params)
+                rows = cur.fetchall()
+        return [
+            {
+                "id": r[0],
+                "tenant_id": r[1],
+                "user_id": r[2],
+                "session_id": r[3],
+                "type": r[4],
+                "content": json.loads(r[5]) if r[5] else {},
+                "metadata": json.loads(r[6]) if r[6] else {},
+                "retention_policy": r[7],
+            }
+            for r in rows
+        ]