amu-pgvector 0.1.1__tar.gz

amu_pgvector-0.1.1/.gitignore

@@ -0,0 +1,24 @@
+__pycache__/
+*.py[cod]
+*.egg-info/
+.eggs/
+build/
+dist/
+.venv/
+venv/
+.uv/
+.uv-cache/
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.coverage
+htmlcov/
+.env
+.env.*
+!.env.example
+*.log
+.DS_Store
+.claude/
+node_modules/
+.idea/
+.vscode/

amu_pgvector-0.1.1/PKG-INFO

@@ -0,0 +1,38 @@
+Metadata-Version: 2.5
+Name: amu-pgvector
+Version: 0.1.1
+Summary: Reference implementation of Lineage-Aware Memory Governance on PostgreSQL + pgvector: Postgres RLS enforces S(a) subset P(d), not application code.
+Project-URL: Homepage, https://github.com/sangaraju1988/amu-pgvector
+Project-URL: Repository, https://github.com/sangaraju1988/amu-pgvector
+Project-URL: Paper, https://doi.org/10.1109/ACCESS.2026.3730363
+Author: Venkata Sangaraju
+License: MIT
+Keywords: agent-memory,governance,lineage,pgvector,postgres,rag,rls,row-level-security
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Database
+Classifier: Topic :: Security
+Requires-Python: >=3.10
+Requires-Dist: amu-governance==0.1.2
+Requires-Dist: pgvector>=0.3
+Requires-Dist: psycopg[binary]>=3.1
+Provides-Extra: async
+Requires-Dist: psycopg[binary,pool]>=3.1; extra == 'async'
+Provides-Extra: mcp
+Requires-Dist: mcp>=2.0; extra == 'mcp'
+Provides-Extra: st
+Requires-Dist: sentence-transformers>=3.0; extra == 'st'
+Description-Content-Type: text/markdown
+
+# amu-pgvector
+
+Python client for [amu-pgvector](https://github.com/sangaraju1988/amu-pgvector) — see the repository root README for the full quickstart, security model, and benchmarks.
+
+The `[mcp]` extra (`pip install amu-pgvector[mcp]`) adds an `amu-pgvector-mcp` console script: an MCP server exposing `amu_search`, `amu_record` and `amu_check_conflict` as lineage-gated tools. See `submissions/mcp-registry.md` in the repository for registry publishing details.
+
+<!-- mcp-name: io.github.sangaraju1988/amu-pgvector -->
+

amu_pgvector-0.1.1/README.md

@@ -0,0 +1,8 @@
+# amu-pgvector
+
+Python client for [amu-pgvector](https://github.com/sangaraju1988/amu-pgvector) — see the repository root README for the full quickstart, security model, and benchmarks.
+
+The `[mcp]` extra (`pip install amu-pgvector[mcp]`) adds an `amu-pgvector-mcp` console script: an MCP server exposing `amu_search`, `amu_record` and `amu_check_conflict` as lineage-gated tools. See `submissions/mcp-registry.md` in the repository for registry publishing details.
+
+<!-- mcp-name: io.github.sangaraju1988/amu-pgvector -->
+

amu_pgvector-0.1.1/pyproject.toml

@@ -0,0 +1,44 @@
+[project]
+name = "amu-pgvector"
+version = "0.1.1"
+description = "Reference implementation of Lineage-Aware Memory Governance on PostgreSQL + pgvector: Postgres RLS enforces S(a) subset P(d), not application code."
+readme = "README.md"
+license = { text = "MIT" }
+requires-python = ">=3.10"
+authors = [{ name = "Venkata Sangaraju" }]
+keywords = ["postgres", "pgvector", "rls", "row-level-security", "agent-memory", "lineage", "governance", "rag"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Database",
+    "Topic :: Security",
+]
+dependencies = [
+    "amu-governance==0.1.2",
+    "psycopg[binary]>=3.1",
+    "pgvector>=0.3",
+]
+
+[project.optional-dependencies]
+mcp = ["mcp>=2.0"]
+st = ["sentence-transformers>=3.0"]
+async = ["psycopg[binary,pool]>=3.1"]
+
+[project.scripts]
+amu-pgvector-mcp = "amu_pgvector.mcp_server:main"
+
+[project.urls]
+Homepage = "https://github.com/sangaraju1988/amu-pgvector"
+Repository = "https://github.com/sangaraju1988/amu-pgvector"
+Paper = "https://doi.org/10.1109/ACCESS.2026.3730363"
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/amu_pgvector"]

amu_pgvector-0.1.1/src/amu_pgvector/__init__.py

@@ -0,0 +1,14 @@
+"""amu-pgvector: Lineage-Aware Memory Governance on PostgreSQL + pgvector.
+
+Postgres row-level security enforces S(a) subset P(d) at the row level, so
+that reuse of a cached analytical result (an Analytical Memory Unit, AMU) is
+gated on its derivation lineage. See sql/amu_pgvector.sql for the enforced
+schema and docs/design.md for how each paper concept maps to a SQL object.
+"""
+
+from amu_pgvector.models import Conflict, RecordResult, SearchResult
+from amu_pgvector.store import AMUStore
+
+__version__ = "0.1.1"
+
+__all__ = ["AMUStore", "Conflict", "RecordResult", "SearchResult"]

amu_pgvector-0.1.1/src/amu_pgvector/embeddings.py

@@ -0,0 +1,62 @@
+"""Pluggable embedding functions.
+
+An embed_fn is just Callable[[str], Sequence[float]]. `fake_embedder` is
+deterministic and needs no network or model download -- it's what the test
+suite and the quickstart use by default. `sentence_transformer_embedder`
+needs the optional `[st]` extra; the import is deferred so the base package
+never requires sentence-transformers/torch.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import struct
+from collections.abc import Callable
+
+
+def fake_embedder(dim: int = 1536) -> Callable[[str], list[float]]:
+    """A deterministic, hash-based embedder: same text always maps to the
+    same vector, semantically meaningless but stable and dependency-free.
+    Two different strings produce different vectors; that's all the test
+    suite (and a quickstart with nothing else installed) needs."""
+
+    def _embed(text: str) -> list[float]:
+        values: list[float] = []
+        counter = 0
+        while len(values) < dim:
+            digest = hashlib.sha256(f"{text}\0{counter}".encode()).digest()
+            for i in range(0, len(digest) - 3, 4):
+                if len(values) >= dim:
+                    break
+                (as_uint,) = struct.unpack(">I", digest[i : i + 4])
+                # Map to [-1, 1].
+                values.append((as_uint / 0xFFFFFFFF) * 2 - 1)
+            counter += 1
+        return values
+
+    return _embed
+
+
+def sentence_transformer_embedder(
+    model_name: str = "all-MiniLM-L6-v2",
+) -> Callable[[str], list[float]]:
+    """Requires the `[st]` extra (sentence-transformers). Real semantic
+    embeddings, for examples and production use -- not for the test suite,
+    which must never need a model download."""
+    try:
+        from sentence_transformers import SentenceTransformer
+    except ImportError as exc:  # pragma: no cover - exercised only without [st]
+        raise ImportError(
+            "sentence_transformer_embedder requires the 'st' extra: "
+            "pip install amu-pgvector[st]"
+        ) from exc
+
+    model = SentenceTransformer(model_name)
+
+    def _embed(text: str) -> list[float]:
+        return model.encode(text, normalize_embeddings=True).tolist()
+
+    return _embed
+
+
+__all__ = ["fake_embedder", "sentence_transformer_embedder"]

amu_pgvector-0.1.1/src/amu_pgvector/mcp_server.py

@@ -0,0 +1,150 @@
+"""amu-pgvector-mcp: an MCP server exposing lineage-gated AMU search,
+record, and conflict-check as tools.
+
+Connects with the DSN of the Postgres role it's launched as (the
+AMU_PGVECTOR_DSN env var) -- results are gated by Postgres row-level
+security under THAT role, exactly as everywhere else in this project. If
+launched as an amu_agent_base-only role, amu_record will fail with a
+Postgres permission error rather than silently degrading -- writes need an
+amu_writer-member role. This server never re-implements or bypasses the
+gate; it just runs SQL as whichever role it was given.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+from mcp.server.mcpserver import MCPServer
+
+from amu_pgvector import AMUStore
+from amu_pgvector.embeddings import fake_embedder
+
+DSN_ENV_VAR = "AMU_PGVECTOR_DSN"
+EMBEDDING_DIM_ENV_VAR = "AMU_PGVECTOR_EMBEDDING_DIM"
+
+_GATING_NOTE = (
+    "Results are lineage-gated by Postgres row-level security under this "
+    "server's own Postgres role: an Analytical Memory Unit (AMU) is only "
+    "ever returned if every sensitive column touched by its derivation is "
+    "in this role's permitted set. This tool cannot bypass that gate."
+)
+
+mcp = MCPServer(
+    name="amu-pgvector",
+    instructions=(
+        "Lineage-gated agent memory on PostgreSQL + pgvector (Lineage-Aware "
+        "Memory Governance, IEEE Access DOI 10.1109/ACCESS.2026.3730363). "
+        + _GATING_NOTE
+    ),
+)
+
+_store: AMUStore | None = None
+_embed_fn = None
+
+
+def _get_store() -> AMUStore:
+    global _store
+    if _store is None:
+        dsn = os.environ.get(DSN_ENV_VAR)
+        if not dsn:
+            raise RuntimeError(
+                f"{DSN_ENV_VAR} must be set to the DSN of the Postgres role "
+                "this server should run as."
+            )
+        _store = AMUStore(dsn)
+    return _store
+
+
+def _get_embed_fn():
+    global _embed_fn
+    if _embed_fn is None:
+        try:
+            from amu_pgvector.embeddings import sentence_transformer_embedder
+
+            _embed_fn = sentence_transformer_embedder()
+        except ImportError:
+            dim = int(os.environ.get(EMBEDDING_DIM_ENV_VAR, "1536"))
+            _embed_fn = fake_embedder(dim=dim)
+    return _embed_fn
+
+
+@mcp.tool(
+    description=(
+        "Search lineage-gated Analytical Memory Units (AMUs) by semantic "
+        "similarity to a natural-language query. " + _GATING_NOTE
+    )
+)
+def amu_search(query: str, k: int = 5, metric_name: str | None = None) -> list[dict[str, Any]]:
+    results = _get_store().search(query, k, embed_fn=_get_embed_fn(), metric_name=metric_name)
+    return [
+        {
+            "id": str(r.id),
+            "metric_name": r.metric_name,
+            "description": r.description,
+            "value": r.value,
+            "owner_department": r.owner_department,
+            "definition_hash": r.definition_hash,
+            "distance": r.distance,
+        }
+        for r in results
+    ]
+
+
+@mcp.tool(
+    description=(
+        "Record a new Analytical Memory Unit (AMU). Lineage is extracted "
+        "from the executed `sql` (ground truth), never self-reported. "
+        "Requires an amu_writer-member Postgres role -- other roles get a "
+        "Postgres permission error, not a silent no-op."
+    )
+)
+def amu_record(
+    sql: str,
+    value: dict[str, Any],
+    metric_name: str,
+    description: str,
+    owner_department: str,
+) -> dict[str, Any]:
+    result = _get_store().record(
+        sql,
+        value,
+        metric_name=metric_name,
+        description=description,
+        owner_department=owner_department,
+        embed_fn=_get_embed_fn(),
+    )
+    return {
+        "amu_id": str(result.amu_id),
+        "lineage_status": result.lineage_status,
+        "conflicts": [
+            {"department": c.department, "definition_hash": c.definition_hash}
+            for c in result.conflicts
+        ],
+    }
+
+
+@mcp.tool(
+    description=(
+        "Check whether another department has already recorded a "
+        "conflicting definition of the same metric_name (same name, "
+        "different derivation / definition_hash). Returns only department "
+        "names and hashes -- never values or lineage -- so this cannot "
+        "become a side channel around the gate."
+    )
+)
+def amu_check_conflict(
+    metric_name: str, definition_hash: str, owner_department: str
+) -> list[dict[str, str]]:
+    conflicts = _get_store().check_conflict(metric_name, definition_hash, owner_department)
+    return [
+        {"department": c.department, "definition_hash": c.definition_hash} for c in conflicts
+    ]
+
+
+def main() -> None:
+    mcp.run()
+
+
+if __name__ == "__main__":
+    main()

amu_pgvector-0.1.1/src/amu_pgvector/models.py

@@ -0,0 +1,36 @@
+"""Small result types returned by AMUStore. Not the governance model itself
+-- that's amu_governance.AMU/Lineage/LineageStep, reused as-is."""
+
+from __future__ import annotations
+
+import uuid
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class Conflict:
+    """A pre-existing AMU for the same metric_name with a different
+    definition_hash, owned by a different department. Carries no values or
+    lineage -- amu.check_conflict() never returns them."""
+
+    department: str
+    definition_hash: str
+
+
+@dataclass(frozen=True)
+class RecordResult:
+    amu_id: uuid.UUID
+    lineage_status: str
+    conflicts: list[Conflict]
+
+
+@dataclass(frozen=True)
+class SearchResult:
+    id: uuid.UUID
+    metric_name: str
+    description: str
+    value: dict
+    owner_department: str
+    definition_hash: str
+    distance: float
+    external_id: str | None = None

amu_pgvector-0.1.1/src/amu_pgvector/store.py

@@ -0,0 +1,269 @@
+"""AMUStore: the Python client for amu-pgvector.
+
+A thin wrapper around one Postgres DSN. Which methods actually work depends
+entirely on what Postgres privileges that DSN's role has -- record() needs
+an amu_writer-member role, search() just needs amu_agent_base, and the
+admin_* helpers need the schema owner. Nothing here re-implements or
+weakens the gate: every security-relevant decision (S(a) subset P(d),
+lineage closure, sensitivity tagging) happens in sql/amu_pgvector.sql, not
+in this client.
+"""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable, Sequence
+
+import psycopg
+from amu_governance.sql_lineage import lineage_from_sql
+from pgvector.psycopg import register_vector
+from psycopg.types.json import Json
+
+from amu_pgvector.models import Conflict, RecordResult, SearchResult
+
+EmbedFn = Callable[[str], Sequence[float]]
+
+
+def _lineage_json(lineage) -> dict:
+    return {
+        "steps": [
+            {"table": step.table, "columns_used": list(step.columns_used)}
+            for step in lineage.steps
+        ],
+        "filter_logic": lineage.filter_logic,
+    }
+
+
+def _row_to_search_result(row: tuple) -> SearchResult:
+    id_, external_id, metric_name, description, value, owner_department, definition_hash, distance = row
+    return SearchResult(
+        id=id_,
+        external_id=external_id,
+        metric_name=metric_name,
+        description=description,
+        value=value,
+        owner_department=owner_department,
+        definition_hash=definition_hash,
+        distance=distance,
+    )
+
+
+class AMUStore:
+    def __init__(self, dsn: str, *, sql_dialect: str = "postgres") -> None:
+        self._dsn = dsn
+        self._sql_dialect = sql_dialect
+
+    def _connect(self) -> psycopg.Connection:
+        conn = psycopg.connect(self._dsn)
+        register_vector(conn)
+        return conn
+
+    def check_conflict(
+        self, metric_name: str, definition_hash: str, owner_department: str
+    ) -> list[Conflict]:
+        """Departments with a different definition_hash for the same
+        metric_name -- department/hash only, never values or lineage (see
+        amu.check_conflict() in sql/amu_pgvector.sql)."""
+        with self._connect() as conn, conn.cursor() as cur:
+            cur.execute(
+                "SELECT conflicting_department, other_hash FROM amu.check_conflict(%s, %s, %s)",
+                (metric_name, definition_hash, owner_department),
+            )
+            return [Conflict(dept, other_hash) for dept, other_hash in cur.fetchall()]
+
+    # -- writer path -----------------------------------------------------
+
+    def record(
+        self,
+        sql: str,
+        result,
+        *,
+        metric_name: str,
+        description: str,
+        owner_department: str,
+        embed_fn: EmbedFn,
+        epoch: int | None = None,
+        external_id: str | None = None,
+    ) -> RecordResult:
+        """Extract lineage from the executed `sql` (ground truth, not
+        self-reported by an agent), compute its definition_hash, check for
+        conflicting definitions of the same metric, insert the AMU, and
+        return the new id plus lineage_status plus any conflicts found.
+
+        `result` is the analytical value being cached (must be JSON-
+        serializable); `description` is the natural-language text that gets
+        embedded for retrieval. `external_id`, if given, upserts (an AMU
+        already recorded under that external_id is updated in place rather
+        than duplicated) -- used by langchain_amu.AMUVectorStore to give
+        add_documents(..., ids=[...]) idempotent/mutable semantics; leave it
+        None for a plain new AMU.
+        """
+        lineage = lineage_from_sql(sql, dialect=self._sql_dialect)
+        definition_hash = lineage.definition_hash()
+        embedding = list(embed_fn(description))
+        if epoch is None:
+            epoch = int(time.time())
+
+        conflicts = self.check_conflict(metric_name, definition_hash, owner_department)
+
+        with self._connect() as conn, conn.cursor() as cur:
+            params = (
+                metric_name,
+                description,
+                Json(result),
+                owner_department,
+                epoch,
+                Json(_lineage_json(lineage)),
+                definition_hash,
+                embedding,
+                external_id,
+            )
+            if external_id is None:
+                cur.execute(
+                    """
+                    INSERT INTO amu.memory_units
+                        (metric_name, description, value, owner_department, epoch,
+                         lineage, definition_hash, embedding, external_id)
+                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
+                    RETURNING id, lineage_status
+                    """,
+                    params,
+                )
+            else:
+                cur.execute(
+                    """
+                    INSERT INTO amu.memory_units
+                        (metric_name, description, value, owner_department, epoch,
+                         lineage, definition_hash, embedding, external_id)
+                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
+                    ON CONFLICT (external_id) WHERE external_id IS NOT NULL DO UPDATE SET
+                        metric_name = EXCLUDED.metric_name,
+                        description = EXCLUDED.description,
+                        value = EXCLUDED.value,
+                        owner_department = EXCLUDED.owner_department,
+                        epoch = EXCLUDED.epoch,
+                        lineage = EXCLUDED.lineage,
+                        definition_hash = EXCLUDED.definition_hash,
+                        embedding = EXCLUDED.embedding
+                    RETURNING id, lineage_status
+                    """,
+                    params,
+                )
+            amu_id, lineage_status = cur.fetchone()
+            conn.commit()
+
+        return RecordResult(amu_id=amu_id, lineage_status=lineage_status, conflicts=conflicts)
+
+    def delete_by_external_ids(self, external_ids: Sequence[str]) -> None:
+        with self._connect() as conn, conn.cursor() as cur:
+            cur.execute(
+                "DELETE FROM amu.memory_units WHERE external_id = ANY(%s)",
+                (list(external_ids),),
+            )
+            conn.commit()
+
+    def get_by_external_ids(self, external_ids: Sequence[str]) -> list[SearchResult]:
+        with self._connect() as conn, conn.cursor() as cur:
+            cur.execute(
+                "SELECT id, external_id, metric_name, description, value, "
+                "owner_department, definition_hash, 0.0 "
+                "FROM amu.memory_units WHERE external_id = ANY(%s)",
+                (list(external_ids),),
+            )
+            return [_row_to_search_result(row) for row in cur.fetchall()]
+
+    # -- agent path (RLS applies under this DSN's own role) --------------
+
+    def search(
+        self,
+        query: str,
+        k: int,
+        embed_fn: EmbedFn,
+        *,
+        metric_name: str | None = None,
+    ) -> list[SearchResult]:
+        embedding = list(embed_fn(query))
+        with self._connect() as conn, conn.cursor() as cur:
+            cur.execute(
+                "SELECT id, external_id, metric_name, description, value, "
+                "owner_department, definition_hash, distance "
+                "FROM amu.search(%s::vector, %s, %s)",
+                (embedding, k, metric_name),
+            )
+            return [_row_to_search_result(row) for row in cur.fetchall()]
+
+    def search_with_vectors(
+        self,
+        embedding: Sequence[float],
+        fetch_k: int,
+        *,
+        metric_name: str | None = None,
+    ) -> list[tuple[SearchResult, list[float]]]:
+        """Like search(), but also returns each candidate's raw embedding --
+        for max_marginal_relevance re-ranking (langchain_amu). Not part of
+        amu.search()'s own return shape (kept lean for the common case);
+        this is a plain SELECT on amu.memory_units instead, so it's still
+        fully RLS-gated under the caller's own role."""
+        embedding = list(embedding)
+        with self._connect() as conn, conn.cursor() as cur:
+            cur.execute(
+                "SELECT id, external_id, metric_name, description, value, owner_department, "
+                "definition_hash, embedding <=> %s::vector AS distance, embedding "
+                "FROM amu.memory_units "
+                "WHERE %s::text IS NULL OR metric_name = %s "
+                "ORDER BY embedding <=> %s::vector LIMIT %s",
+                (embedding, metric_name, metric_name, embedding, fetch_k),
+            )
+            rows = cur.fetchall()
+        return [(_row_to_search_result(row[:-1]), row[-1].to_list()) for row in rows]
+
+    # -- admin helpers (need schema-owner privilege) ----------------------
+
+    def register_sensitive_column(self, column_id: str) -> None:
+        with self._connect() as conn, conn.cursor() as cur:
+            cur.execute(
+                "INSERT INTO amu.sensitive_columns (column_id) VALUES (%s) "
+                "ON CONFLICT DO NOTHING",
+                (column_id,),
+            )
+            conn.commit()
+
+    def grant_department_permission(self, department: str, column_id: str) -> None:
+        with self._connect() as conn, conn.cursor() as cur:
+            cur.execute(
+                "INSERT INTO amu.department_permissions (department, column_id) "
+                "VALUES (%s, %s) ON CONFLICT DO NOTHING",
+                (department, column_id),
+            )
+            conn.commit()
+
+    def create_agent_role(self, role_name: str, department: str, password: str) -> None:
+        """Create a LOGIN role that's a member of amu_agent_base and mapped
+        to `department` in amu.role_departments -- the identity mechanism
+        strong mode relies on (session_user, not anything the client sets).
+        """
+        from psycopg import sql as pg_sql
+
+        with self._connect() as conn, conn.cursor() as cur:
+            cur.execute(
+                pg_sql.SQL("CREATE ROLE {} LOGIN PASSWORD {} IN ROLE amu_agent_base").format(
+                    pg_sql.Identifier(role_name), pg_sql.Literal(password)
+                )
+            )
+            cur.execute(
+                "INSERT INTO amu.role_departments (role_name, department) VALUES (%s, %s)",
+                (role_name, department),
+            )
+            conn.commit()
+
+    def register_materialization_edge(
+        self, derived_table: str, source_table: str, column_map: dict[str, str | list[str]]
+    ) -> None:
+        with self._connect() as conn, conn.cursor() as cur:
+            cur.execute(
+                "INSERT INTO amu.materialization_edges "
+                "(derived_table, source_table, column_map) VALUES (%s, %s, %s) "
+                "ON CONFLICT (derived_table, source_table) DO UPDATE SET column_map = EXCLUDED.column_map",
+                (derived_table, source_table, Json(column_map)),
+            )
+            conn.commit()