hermes-memory-pgvector 0.3.0__py3-none-any.whl

pgvector/store.py

@@ -0,0 +1,507 @@
+"""store.py — Postgres ops for the pgvector memory plugin.
+
+Wraps psycopg3 + psycopg_pool. Mirrors hermes-agent's native built-in
+memory model (`memory` tool's add/replace/remove on targets 'memory' /
+'user') into a single Postgres table with embeddings.
+
+Uses a small ConnectionPool because the plugin is touched from two
+threads at runtime: the agent thread (for prefetch / recall_memory /
+ensure_schema / health) and the async-writer drain thread (for the
+mirrored INSERTs / UPDATEs / DELETEs). Pooling beats short-lived
+connections under that two-thread pattern without adding much
+complexity.
+
+No SQLAlchemy, no LLM-mediated workers, no deriver loops.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import threading
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import psycopg
+from psycopg.rows import dict_row
+from psycopg_pool import ConnectionPool
+
+from .embed import to_pgvector_literal
+
+logger = logging.getLogger(__name__)
+
+
+class MemoryStore:
+    """Postgres-backed mirror of hermes-agent's built-in memory entries."""
+
+    def __init__(
+        self,
+        dsn: str,
+        *,
+        min_size: int = 1,
+        max_size: int = 4,
+        timeout: float = 5.0,
+    ):
+        """Open a lazily-initialized ConnectionPool.
+
+        min_size=1 keeps a single warm connection — enough for the agent
+        thread under low traffic. max_size=4 lets the writer thread + a
+        burst of recall calls overlap. Pool is opened on first use (or
+        explicit `open()`) and closed via `close()`.
+        """
+        self._dsn = dsn
+        self._lock = threading.Lock()
+        self._pool: Optional[ConnectionPool] = None
+        self._min_size = min_size
+        self._max_size = max_size
+        self._timeout = timeout
+
+    # -- Pool lifecycle ------------------------------------------------------
+
+    def _get_pool(self) -> ConnectionPool:
+        """Return the live pool, constructing it on first call. Thread-safe."""
+        if self._pool is not None:
+            return self._pool
+        with self._lock:
+            if self._pool is None:
+                self._pool = ConnectionPool(
+                    conninfo=self._dsn,
+                    min_size=self._min_size,
+                    max_size=self._max_size,
+                    timeout=self._timeout,
+                    open=True,
+                    name="pgvector-memory",
+                )
+        return self._pool
+
+    def close(self) -> None:
+        """Close the connection pool. Idempotent."""
+        with self._lock:
+            if self._pool is not None:
+                try:
+                    self._pool.close()
+                except Exception as exc:  # noqa: BLE001
+                    logger.debug("pgvector pool close: %s", exc)
+                finally:
+                    self._pool = None
+
+    # -- Schema --------------------------------------------------------------
+
+    class SchemaNotApplied(RuntimeError):
+        """Raised when memory_entries does not exist in the target DB."""
+
+    def ensure_schema(self) -> None:
+        """Verify the schema is in place. Does NOT run DDL.
+
+        The migration (migrations/001_schema.sql) is admin-only — it
+        runs `CREATE EXTENSION vector` which requires superuser, and
+        creates the table + indexes which then end up owned by the
+        admin role. The plugin's runtime user (hermes) only has
+        SELECT/INSERT/UPDATE/DELETE on the existing schema, and that's
+        the right separation: DDL at install time, DML at run time.
+
+        Operators apply the migration once via:
+            sudo -u postgres psql -d hermes_memory -f migrations/001_schema.sql
+        """
+        with self._get_pool().connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute("SELECT to_regclass('memory_entries')")
+                if cur.fetchone()[0] is None:
+                    raise self.SchemaNotApplied(
+                        "memory_entries table missing. Apply the migration as DB admin: "
+                        "psql -d <dbname> -f plugins/memory/pgvector/migrations/001_schema.sql"
+                    )
+
+    def apply_migration_as_admin(self, *, admin_dsn: str) -> None:
+        """One-shot admin path: run the full migration with privileged creds.
+
+        Bypasses the runtime pool — opens a fresh autocommit connection
+        with admin_dsn (typically `user=postgres host=/var/run/postgresql`)
+        so CREATE EXTENSION + CREATE TABLE + CREATE INDEX all succeed.
+        Idempotent: re-running on an already-migrated DB is a no-op.
+        """
+        sql_path = Path(__file__).parent / "migrations" / "001_schema.sql"
+        sql = sql_path.read_text(encoding="utf-8")
+        with psycopg.connect(admin_dsn, autocommit=True) as conn:
+            with conn.cursor() as cur:
+                cur.execute(sql)
+
+    # -- Built-in memory mirror (called by on_memory_write) ------------------
+
+    def add(
+        self,
+        *,
+        agent_identity: str,
+        target: str,
+        content: str,
+        embedding: Optional[List[float]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> Optional[int]:
+        """Insert a memory entry. Returns row id, or None if duplicate (no-op).
+
+        Matches the built-in tool's "reject exact duplicate" semantics via
+        the (agent_identity, target, content) unique constraint + ON CONFLICT.
+        """
+        meta_json = json.dumps(metadata or {})
+        vec_literal = to_pgvector_literal(embedding) if embedding is not None else None
+
+        with self._get_pool().connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    """
+                    INSERT INTO memory_entries
+                        (agent_identity, target, content, embedding, metadata)
+                    VALUES (%s, %s, %s, %s::vector, %s::jsonb)
+                    ON CONFLICT (agent_identity, target, content) DO NOTHING
+                    RETURNING id
+                    """,
+                    (agent_identity, target, content, vec_literal, meta_json),
+                )
+                row = cur.fetchone()
+                conn.commit()
+                return int(row[0]) if row else None
+
+    def replace(
+        self,
+        *,
+        agent_identity: str,
+        target: str,
+        old_text: str,
+        new_content: str,
+        new_embedding: Optional[List[float]] = None,
+    ) -> int:
+        """Update entries in (agent_identity, target) where content contains old_text.
+
+        Matches built-in semantics — old_text is a substring match. Returns
+        the number of rows updated (built-in updates the FIRST match; we
+        update all matches in the same scope for safety).
+        """
+        vec_literal = (
+            to_pgvector_literal(new_embedding) if new_embedding is not None else None
+        )
+        with self._get_pool().connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    """
+                    UPDATE memory_entries
+                       SET content    = %s,
+                           embedding  = %s::vector,
+                           updated_at = now()
+                     WHERE agent_identity = %s
+                       AND target = %s
+                       AND content LIKE %s
+                    """,
+                    (new_content, vec_literal, agent_identity, target, f"%{old_text}%"),
+                )
+                updated = cur.rowcount
+                conn.commit()
+                return int(updated)
+
+    def remove(
+        self,
+        *,
+        agent_identity: str,
+        target: str,
+        old_text: str,
+    ) -> int:
+        """Delete entries in (agent_identity, target) matching old_text substring.
+
+        Returns the number of rows deleted.
+        """
+        with self._get_pool().connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    """
+                    DELETE FROM memory_entries
+                     WHERE agent_identity = %s
+                       AND target = %s
+                       AND content LIKE %s
+                    """,
+                    (agent_identity, target, f"%{old_text}%"),
+                )
+                deleted = cur.rowcount
+                conn.commit()
+                return int(deleted)
+
+    # -- Reads ---------------------------------------------------------------
+
+    def list_entries(
+        self,
+        *,
+        agent_identity: str,
+        target: Optional[str] = None,
+        limit: int = 100,
+    ) -> List[Dict[str, Any]]:
+        """List entries in an agent's scope. If target is None, both stores."""
+        params: List[Any] = [agent_identity]
+        target_clause = ""
+        if target:
+            target_clause = "AND target = %s"
+            params.append(target)
+        params.append(limit)
+
+        with self._get_pool().connection() as conn:
+            with conn.cursor(row_factory=dict_row) as cur:
+                cur.execute(
+                    f"""
+                    SELECT id, agent_identity, target, content, created_at, updated_at, metadata
+                    FROM memory_entries
+                    WHERE agent_identity = %s
+                    {target_clause}
+                    ORDER BY updated_at DESC
+                    LIMIT %s
+                    """,
+                    params,
+                )
+                return list(cur.fetchall())
+
+    def search(
+        self,
+        *,
+        query_embedding: List[float],
+        agent_identity: Optional[str] = None,
+        target: Optional[str] = None,
+        limit: int = 5,
+        min_similarity: float = 0.0,
+    ) -> List[Dict[str, Any]]:
+        """Semantic recall via cosine distance.
+
+        agent_identity=None → search across ALL agents (cross-theme recall).
+        target=None → search both 'memory' and 'user'.
+        Returns rows with `score` = 1 - cosine_distance ∈ [0, 1].
+        """
+        vec_literal = to_pgvector_literal(query_embedding)
+        clauses: List[str] = []
+        params: List[Any] = []
+        if agent_identity:
+            clauses.append("agent_identity = %s")
+            params.append(agent_identity)
+        if target:
+            clauses.append("target = %s")
+            params.append(target)
+        where = ("WHERE " + " AND ".join(clauses)) if clauses else ""
+
+        with self._get_pool().connection() as conn:
+            with conn.cursor(row_factory=dict_row) as cur:
+                cur.execute(
+                    f"""
+                    SELECT id, agent_identity, target, content, created_at,
+                           updated_at, metadata,
+                           1 - (embedding <=> %s::vector) AS score
+                    FROM memory_entries
+                    {where}
+                    ORDER BY embedding <=> %s::vector
+                    LIMIT %s
+                    """,
+                    [vec_literal, *params, vec_literal, limit],
+                )
+                rows = list(cur.fetchall())
+
+        if min_similarity > 0:
+            rows = [r for r in rows if (r.get("score") or 0) >= min_similarity]
+        return rows
+
+    # -- Bulk import from MEMORY.md / USER.md (v0.1.1) ----------------------
+
+    # Matches tools/memory_tool.py:ENTRY_DELIMITER. Keep in sync if upstream
+    # ever changes it (currently stable; been "\n§\n" since the tool shipped).
+    ENTRY_DELIMITER = "\n§\n"
+
+    def bulk_upsert_md(
+        self,
+        *,
+        agent_identity: str,
+        target: str,
+        file_path: "Path | str",
+        embed_fn,
+    ) -> Dict[str, int]:
+        """Parse a MEMORY.md / USER.md file and upsert each entry.
+
+        Idempotent + cheap on re-run: we SELECT the existing content set
+        for (agent_identity, target) once, then only embed + INSERT new
+        entries. So initial install embeds everything; subsequent inits
+        with no MD changes do zero embed calls.
+
+        embed_fn is a callable taking a string and returning a 768-dim
+        list (or raising — we catch and store text-only). Wired by the
+        caller so the plugin can pass its `embed()` with the configured
+        base_url + model.
+
+        Returns: {'parsed': N, 'inserted': M, 'skipped': K} where N=M+K.
+        """
+        from pathlib import Path as _Path
+        p = _Path(file_path)
+        if not p.exists():
+            return {"parsed": 0, "inserted": 0, "skipped": 0}
+
+        raw = p.read_text(encoding="utf-8", errors="replace")
+        entries = [e.strip() for e in raw.split(self.ENTRY_DELIMITER) if e.strip()]
+        if not entries:
+            return {"parsed": 0, "inserted": 0, "skipped": 0}
+
+        # Single bulk SELECT of existing content for this scope. Beats N+1
+        # by a wide margin and keeps re-init nearly free.
+        with self._get_pool().connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    "SELECT content FROM memory_entries WHERE agent_identity = %s AND target = %s",
+                    (agent_identity, target),
+                )
+                existing = {row[0] for row in cur.fetchall()}
+
+        inserted = 0
+        skipped = 0
+        for entry in entries:
+            if entry in existing:
+                skipped += 1
+                continue
+            vec = None
+            try:
+                vec = embed_fn(entry) if embed_fn else None
+            except Exception:  # noqa: BLE001 — fail-soft on bulk embed
+                vec = None
+            row_id = self.add(
+                agent_identity=agent_identity,
+                target=target,
+                content=entry,
+                embedding=vec,
+                metadata={"source": "bulk_import", "file": str(p)},
+            )
+            if row_id is not None:
+                inserted += 1
+            else:
+                # Lost a race with another writer that inserted the same row.
+                skipped += 1
+        return {"parsed": len(entries), "inserted": inserted, "skipped": skipped}
+
+    # -- Conversation turns (v0.2) ------------------------------------------
+
+    def append_turn(
+        self,
+        *,
+        session_id: str,
+        agent_identity: str,
+        role: str,
+        content: str,
+        embedding: Optional[List[float]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> int:
+        """Insert one chat turn. Returns row id.
+
+        No dedup (turns are inherently time-ordered events — same content
+        twice is two distinct turns, even verbatim).
+        """
+        meta_json = json.dumps(metadata or {})
+        vec_literal = to_pgvector_literal(embedding) if embedding is not None else None
+
+        with self._get_pool().connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(
+                    """
+                    INSERT INTO conversations
+                        (session_id, agent_identity, role, content, embedding, metadata)
+                    VALUES (%s, %s, %s, %s, %s::vector, %s::jsonb)
+                    RETURNING id
+                    """,
+                    (session_id, agent_identity, role, content, vec_literal, meta_json),
+                )
+                row = cur.fetchone()
+                conn.commit()
+                return int(row[0])
+
+    def search_turns(
+        self,
+        *,
+        query_embedding: List[float],
+        agent_identity: Optional[str] = None,
+        session_id: Optional[str] = None,
+        limit: int = 5,
+        min_similarity: float = 0.0,
+    ) -> List[Dict[str, Any]]:
+        """Semantic recall over conversation turns. Same shape as `search()`."""
+        vec_literal = to_pgvector_literal(query_embedding)
+        clauses: List[str] = []
+        params: List[Any] = []
+        if agent_identity:
+            clauses.append("agent_identity = %s")
+            params.append(agent_identity)
+        if session_id:
+            clauses.append("session_id = %s")
+            params.append(session_id)
+        where = ("WHERE " + " AND ".join(clauses)) if clauses else ""
+
+        with self._get_pool().connection() as conn:
+            with conn.cursor(row_factory=dict_row) as cur:
+                cur.execute(
+                    f"""
+                    SELECT id, session_id, agent_identity, role, content, ts, metadata,
+                           1 - (embedding <=> %s::vector) AS score
+                    FROM conversations
+                    {where}
+                    ORDER BY embedding <=> %s::vector
+                    LIMIT %s
+                    """,
+                    [vec_literal, *params, vec_literal, limit],
+                )
+                rows = list(cur.fetchall())
+
+        if min_similarity > 0:
+            rows = [r for r in rows if (r.get("score") or 0) >= min_similarity]
+        return rows
+
+    # -- Maintenance ---------------------------------------------------------
+
+    def count_turns(
+        self,
+        *,
+        agent_identity: Optional[str] = None,
+        session_id: Optional[str] = None,
+    ) -> int:
+        clauses: List[str] = []
+        params: List[Any] = []
+        if agent_identity:
+            clauses.append("agent_identity = %s")
+            params.append(agent_identity)
+        if session_id:
+            clauses.append("session_id = %s")
+            params.append(session_id)
+        where = ("WHERE " + " AND ".join(clauses)) if clauses else ""
+        with self._get_pool().connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(f"SELECT COUNT(*) FROM conversations {where}", params)
+                return int(cur.fetchone()[0])
+
+    def count(
+        self,
+        *,
+        agent_identity: Optional[str] = None,
+        target: Optional[str] = None,
+    ) -> int:
+        clauses: List[str] = []
+        params: List[Any] = []
+        if agent_identity:
+            clauses.append("agent_identity = %s")
+            params.append(agent_identity)
+        if target:
+            clauses.append("target = %s")
+            params.append(target)
+        where = ("WHERE " + " AND ".join(clauses)) if clauses else ""
+
+        with self._get_pool().connection() as conn:
+            with conn.cursor() as cur:
+                cur.execute(f"SELECT COUNT(*) FROM memory_entries {where}", params)
+                return int(cur.fetchone()[0])
+
+    def health(self) -> Dict[str, Any]:
+        """Liveness probe — pool reachable + table exists. Never raises."""
+        try:
+            with self._get_pool().connection(timeout=3.0) as conn:
+                with conn.cursor() as cur:
+                    cur.execute("SELECT to_regclass('memory_entries') IS NOT NULL")
+                    has_table = bool(cur.fetchone()[0])
+                    if not has_table:
+                        return {"ok": False, "error": "memory_entries table missing", "row_count": 0}
+                    cur.execute("SELECT COUNT(*) FROM memory_entries")
+                    return {"ok": True, "error": "", "row_count": int(cur.fetchone()[0])}
+        except Exception as exc:  # noqa: BLE001
+            return {"ok": False, "error": str(exc)[:200], "row_count": 0}

pgvector/writer.py

@@ -0,0 +1,170 @@
+"""writer.py — async write queue for the pgvector memory plugin.
+
+Decouples `on_memory_write` from the (potentially slow) embed + INSERT
+path. Built-in memory tool fires the hook → we enqueue → background
+thread drains, embeds, writes. The hook returns in microseconds, never
+blocks the agent loop on a slow embed endpoint.
+
+Failure modes we deliberately swallow:
+  • queue full              → drop write, log once, set a tripwire flag
+  • embed endpoint slow/down → write goes through with embedding=NULL,
+                              row is still recoverable, recall just
+                              misses it until a future backfill
+  • DB temporarily down     → drop write, log once
+  • background thread crash → log + auto-restart on next enqueue
+
+The whole point is to avoid the Honcho-style retry-storm meltdown.
+"""
+
+from __future__ import annotations
+
+import logging
+import queue
+import threading
+from typing import Any, Callable, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+
+# Pending-write payload — kept small so the queue stays bounded.
+class _PendingWrite:
+    __slots__ = ("action", "agent_identity", "target", "content", "extra", "metadata")
+
+    def __init__(
+        self,
+        *,
+        action: str,
+        agent_identity: str,
+        target: str,
+        content: str,
+        extra: Optional[Dict[str, Any]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ):
+        self.action = action
+        self.agent_identity = agent_identity
+        self.target = target
+        self.content = content
+        self.extra = extra or {}
+        self.metadata = metadata or {}
+
+
+class AsyncWriter:
+    """Bounded write queue with a background drain thread.
+
+    Caller responsibilities:
+      - Construct with `worker_fn`, a callable that takes a _PendingWrite
+        and performs the actual embed + DB write. worker_fn MUST NOT raise
+        — it should swallow + log its own failures.
+      - Call `enqueue(...)` from the agent thread; returns instantly.
+      - Call `shutdown(timeout=5)` on plugin teardown to drain gracefully.
+    """
+
+    def __init__(self, worker_fn: Callable[[_PendingWrite], None], *, maxsize: int = 256):
+        self._worker_fn = worker_fn
+        self._queue: "queue.Queue[Optional[_PendingWrite]]" = queue.Queue(maxsize=maxsize)
+        self._thread: Optional[threading.Thread] = None
+        self._stop = threading.Event()
+        self._dropped = 0
+        self._dropped_warned = False
+        self._lock = threading.Lock()
+
+    # -- Public API ----------------------------------------------------------
+
+    def enqueue(
+        self,
+        *,
+        action: str,
+        agent_identity: str,
+        target: str,
+        content: str,
+        extra: Optional[Dict[str, Any]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> bool:
+        """Enqueue a write. Returns True on accept, False on drop (queue full).
+
+        Lazily starts the drain thread on first enqueue. Restarts the
+        thread if it died (e.g., from an unforeseen exception in worker_fn).
+        """
+        self._ensure_thread()
+        item = _PendingWrite(
+            action=action,
+            agent_identity=agent_identity,
+            target=target,
+            content=content,
+            extra=extra,
+            metadata=metadata,
+        )
+        try:
+            self._queue.put_nowait(item)
+            return True
+        except queue.Full:
+            with self._lock:
+                self._dropped += 1
+                if not self._dropped_warned:
+                    logger.warning(
+                        "pgvector writer queue full (maxsize=%d); dropping writes",
+                        self._queue.maxsize,
+                    )
+                    self._dropped_warned = True
+            return False
+
+    def shutdown(self, timeout: float = 5.0) -> None:
+        """Signal stop + wait for the queue to drain (up to `timeout` seconds)."""
+        if not self._thread or not self._thread.is_alive():
+            return
+        # Sentinel wakes the thread even if the queue is empty.
+        try:
+            self._queue.put_nowait(None)
+        except queue.Full:
+            self._stop.set()
+        self._thread.join(timeout=timeout)
+        if self._thread.is_alive():
+            logger.warning("pgvector writer thread did not drain within %.1fs", timeout)
+
+    def stats(self) -> Dict[str, Any]:
+        return {
+            "queue_size": self._queue.qsize(),
+            "queue_max": self._queue.maxsize,
+            "dropped_total": self._dropped,
+            "thread_alive": bool(self._thread and self._thread.is_alive()),
+        }
+
+    # -- Internals -----------------------------------------------------------
+
+    def _ensure_thread(self) -> None:
+        with self._lock:
+            if self._thread and self._thread.is_alive():
+                return
+            self._stop.clear()
+            self._thread = threading.Thread(
+                target=self._run,
+                name="pgvector-writer",
+                daemon=True,
+            )
+            self._thread.start()
+
+    def _run(self) -> None:
+        while not self._stop.is_set():
+            try:
+                item = self._queue.get(timeout=0.5)
+            except queue.Empty:
+                continue
+            if item is None:  # shutdown sentinel
+                self._queue.task_done()
+                break
+            try:
+                self._worker_fn(item)
+            except Exception as exc:  # noqa: BLE001
+                # worker_fn is supposed to swallow its own errors. If
+                # something slipped through, log and keep going — never
+                # let the drain thread die over one bad write.
+                logger.warning(
+                    "pgvector writer worker raised %s on (%s/%s/%s): %s",
+                    type(exc).__name__,
+                    item.action,
+                    item.agent_identity,
+                    item.target,
+                    str(exc)[:200],
+                )
+            finally:
+                self._queue.task_done()