astrocyte-pgvector 0.1.0__py3-none-any.whl

astrocyte_pgvector/__init__.py

@@ -0,0 +1,5 @@
+"""PostgreSQL + pgvector adapter for Astrocyte Tier 1."""
+
+from astrocyte_pgvector.store import PgVectorStore
+
+__all__ = ["PgVectorStore"]

astrocyte_pgvector/store.py

@@ -0,0 +1,280 @@
+"""VectorStore backed by PostgreSQL with the pgvector extension."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import re
+from datetime import datetime
+from typing import Any, ClassVar
+
+import psycopg
+from pgvector.psycopg import register_vector_async
+from psycopg.rows import dict_row
+from psycopg.types.json import Json
+from psycopg_pool import AsyncConnectionPool
+
+from astrocyte.types import HealthStatus, VectorFilters, VectorHit, VectorItem
+
+_TABLE_SAFE = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$")
+
+
+def _sanitize_table(name: str) -> str:
+    if not _TABLE_SAFE.match(name):
+        raise ValueError(f"Invalid table name: {name!r}")
+    return name
+
+
+class PgVectorStore:
+    """Tier 1 vector store using `pgvector` cosine distance search."""
+
+    SPI_VERSION: ClassVar[int] = 1
+
+    def __init__(
+        self,
+        dsn: str | None = None,
+        table_name: str = "astrocyte_vectors",
+        embedding_dimensions: int = 128,
+        bootstrap_schema: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        self._dsn = dsn or os.environ.get("DATABASE_URL") or os.environ.get("ASTROCYTES_PG_DSN")
+        if not self._dsn:
+            raise ValueError(
+                "PgVectorStore requires `dsn` in vector_store_config or DATABASE_URL / ASTROCYTES_PG_DSN",
+            )
+        self._table = _sanitize_table(table_name)
+        self._dim = int(embedding_dimensions)
+        if self._dim < 1:
+            raise ValueError("embedding_dimensions must be >= 1")
+        self._bootstrap_schema = bool(bootstrap_schema)
+        self._pool: AsyncConnectionPool | None = None
+        self._pool_lock = asyncio.Lock()
+        # When migrations own DDL, skip in-app CREATE TABLE / indexes.
+        self._schema_ready = not self._bootstrap_schema
+        self._schema_lock = asyncio.Lock()
+
+    async def _ensure_pool(self) -> AsyncConnectionPool:
+        async with self._pool_lock:
+            if self._pool is None:
+
+                async def configure(conn: psycopg.AsyncConnection) -> None:
+                    await conn.execute("SELECT 1")
+                    # register_vector_async needs the `vector` type. Skip until pgvector exists (quick path:
+                    # /health can run before in-app DDL; runbook path: migrations already created the extension).
+                    async with conn.cursor() as cur:
+                        await cur.execute("SELECT 1 FROM pg_extension WHERE extname = 'vector'")
+                        ext_present = await cur.fetchone()
+                    if ext_present:
+                        await register_vector_async(conn)
+                    await conn.commit()
+
+                self._pool = AsyncConnectionPool(
+                    conninfo=self._dsn,
+                    configure=configure,
+                    open=False,
+                    min_size=1,
+                    max_size=10,
+                    kwargs={"connect_timeout": 10},
+                )
+                await self._pool.open()
+            return self._pool
+
+    async def _ensure_schema(self, pool: AsyncConnectionPool) -> None:
+        async with self._schema_lock:
+            if self._schema_ready:
+                return
+            async with pool.connection() as conn:
+                await conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
+                await conn.execute(
+                    f"""
+                    CREATE TABLE IF NOT EXISTS {self._table} (
+                        id TEXT PRIMARY KEY,
+                        bank_id TEXT NOT NULL,
+                        embedding vector({self._dim}) NOT NULL,
+                        text TEXT NOT NULL,
+                        metadata JSONB,
+                        tags TEXT[],
+                        fact_type TEXT,
+                        occurred_at TIMESTAMPTZ,
+                        memory_layer TEXT
+                    )
+                    """
+                )
+                await conn.execute(
+                    f"CREATE INDEX IF NOT EXISTS {self._table}_bank_idx ON {self._table} (bank_id)"
+                )
+                await register_vector_async(conn)
+                await conn.commit()
+            self._schema_ready = True
+
+    async def store_vectors(self, items: list[VectorItem]) -> list[str]:
+        pool = await self._ensure_pool()
+        await self._ensure_schema(pool)
+        stored: list[str] = []
+        async with pool.connection() as conn:
+            async with conn.cursor() as cur:
+                for item in items:
+                    if len(item.vector) != self._dim:
+                        raise ValueError(
+                            f"Vector length {len(item.vector)} != embedding_dimensions {self._dim}",
+                        )
+                    await cur.execute(
+                        f"""
+                        INSERT INTO {self._table}
+                            (id, bank_id, embedding, text, metadata, tags, fact_type, occurred_at, memory_layer)
+                        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
+                        ON CONFLICT (id) DO UPDATE SET
+                            bank_id = EXCLUDED.bank_id,
+                            embedding = EXCLUDED.embedding,
+                            text = EXCLUDED.text,
+                            metadata = EXCLUDED.metadata,
+                            tags = EXCLUDED.tags,
+                            fact_type = EXCLUDED.fact_type,
+                            occurred_at = EXCLUDED.occurred_at,
+                            memory_layer = EXCLUDED.memory_layer
+                        """,
+                        (
+                            item.id,
+                            item.bank_id,
+                            item.vector,
+                            item.text,
+                            Json(item.metadata) if item.metadata is not None else None,
+                            item.tags,
+                            item.fact_type,
+                            item.occurred_at,
+                            item.memory_layer,
+                        ),
+                    )
+                    stored.append(item.id)
+        return stored
+
+    async def search_similar(
+        self,
+        query_vector: list[float],
+        bank_id: str,
+        limit: int = 10,
+        filters: VectorFilters | None = None,
+    ) -> list[VectorHit]:
+        if len(query_vector) != self._dim:
+            raise ValueError(
+                f"Query vector length {len(query_vector)} != embedding_dimensions {self._dim}",
+            )
+        pool = await self._ensure_pool()
+        await self._ensure_schema(pool)
+
+        where = ["bank_id = %s"]
+        params: list[Any] = [query_vector, bank_id]
+        if filters and filters.tags:
+            where.append("tags && %s::text[]")
+            params.append(filters.tags)
+        if filters and filters.fact_types:
+            where.append("fact_type = ANY(%s::text[])")
+            params.append(filters.fact_types)
+        params.extend([query_vector, limit])
+
+        where_sql = " AND ".join(where)
+        # Cosine distance `<=>`; map to a 0–1-ish score via (1 - distance).
+        sql = f"""
+            SELECT id, text, metadata, tags, fact_type, occurred_at, memory_layer,
+                   (1 - (embedding <=> %s::vector))::float AS score
+            FROM {self._table}
+            WHERE {where_sql}
+            ORDER BY embedding <=> %s::vector
+            LIMIT %s
+        """
+
+        async with pool.connection() as conn:
+            async with conn.cursor(row_factory=dict_row) as cur:
+                await cur.execute(sql, params)
+                rows = await cur.fetchall()
+
+        hits: list[VectorHit] = []
+        for row in rows:
+            score = float(row["score"])
+            if score < 0.0:
+                score = 0.0
+            if score > 1.0:
+                score = 1.0
+            md = row["metadata"]
+            if isinstance(md, str):
+                md = json.loads(md)
+            hits.append(
+                VectorHit(
+                    id=row["id"],
+                    text=row["text"],
+                    score=score,
+                    metadata=md,
+                    tags=list(row["tags"]) if row["tags"] else None,
+                    fact_type=row["fact_type"],
+                    occurred_at=row["occurred_at"],
+                    memory_layer=row.get("memory_layer"),
+                )
+            )
+        return hits
+
+    async def list_vectors(
+        self,
+        bank_id: str,
+        offset: int = 0,
+        limit: int = 100,
+    ) -> list[VectorItem]:
+        pool = await self._ensure_pool()
+        await self._ensure_schema(pool)
+        async with pool.connection() as conn:
+            async with conn.cursor(row_factory=dict_row) as cur:
+                await cur.execute(
+                    f"""
+                    SELECT id, bank_id, embedding, text, metadata, tags, fact_type,
+                           occurred_at, memory_layer
+                    FROM {self._table}
+                    WHERE bank_id = %s
+                    ORDER BY id
+                    OFFSET %s LIMIT %s
+                    """,
+                    (bank_id, offset, limit),
+                )
+                rows = await cur.fetchall()
+        items: list[VectorItem] = []
+        for row in rows:
+            md = row["metadata"]
+            if isinstance(md, str):
+                md = json.loads(md)
+            items.append(
+                VectorItem(
+                    id=row["id"],
+                    bank_id=row["bank_id"],
+                    vector=list(row["embedding"]),
+                    text=row["text"],
+                    metadata=md,
+                    tags=list(row["tags"]) if row["tags"] else None,
+                    fact_type=row["fact_type"],
+                    occurred_at=row["occurred_at"],
+                    memory_layer=row.get("memory_layer"),
+                )
+            )
+        return items
+
+    async def delete(self, ids: list[str], bank_id: str) -> int:
+        if not ids:
+            return 0
+        pool = await self._ensure_pool()
+        await self._ensure_schema(pool)
+        async with pool.connection() as conn:
+            async with conn.cursor() as cur:
+                await cur.execute(
+                    f"DELETE FROM {self._table} WHERE bank_id = %s AND id = ANY(%s::text[])",
+                    (bank_id, ids),
+                )
+                return cur.rowcount or 0
+
+    async def health(self) -> HealthStatus:
+        try:
+            pool = await self._ensure_pool()
+            async with pool.connection() as conn:
+                async with conn.cursor() as cur:
+                    await cur.execute("SELECT 1")
+            return HealthStatus(healthy=True, message="pgvector connected")
+        except Exception as e:
+            return HealthStatus(healthy=False, message=f"pgvector unhealthy: {e!s}")

astrocyte_pgvector-0.1.0.dist-info/METADATA

@@ -0,0 +1,119 @@
+Metadata-Version: 2.4
+Name: astrocyte-pgvector
+Version: 0.1.0
+Summary: PostgreSQL pgvector VectorStore adapter for Astrocyte
+License-Expression: Apache-2.0
+Requires-Python: >=3.11
+Requires-Dist: astrocyte<0.2,>=0.1.0
+Requires-Dist: pgvector>=0.4
+Requires-Dist: psycopg-pool>=3.2
+Requires-Dist: psycopg[binary]>=3.1
+Description-Content-Type: text/markdown
+
+# astrocyte-pgvector
+
+**PostgreSQL + [pgvector](https://github.com/pgvector/pgvector)** implementation of the Astrocyte **`VectorStore`** SPI ([`provider-spi.md`](../../docs/_plugins/provider-spi.md)).
+
+## Install
+
+From the monorepo (with `astrocyte` available):
+
+```bash
+cd astrocyte-services-py/astrocyte-pgvector
+uv sync
+# or: pip install -e ../../astrocyte-py && pip install -e .
+```
+
+Entry point name: **`pgvector`** (group `astrocyte.vector_stores`).
+
+## PostgreSQL with Docker
+
+Use the **combined** Compose stack in **[`../docker-compose.yml`](../docker-compose.yml)** (directory **`astrocyte-services-py/`**) to run **Postgres (pgvector) + the reference REST service** together:
+
+```bash
+cd astrocyte-services-py
+docker compose up -d
+```
+
+For **Postgres only** (no HTTP), start only `postgres`:
+
+```bash
+cd astrocyte-services-py
+docker compose up -d postgres
+```
+
+Default DSN from your host (port **5433** maps to Postgres in the compose file):
+
+```text
+postgresql://astrocyte:astrocyte@127.0.0.1:5433/astrocyte
+```
+
+## Schema migrations (production)
+
+DDL is shipped as **plain SQL** under [`migrations/`](migrations/) and applied with **`psql`** via [`scripts/migrate.sh`](scripts/migrate.sh) (no Python migration framework).
+
+```bash
+export DATABASE_URL='postgresql://astrocyte:astrocyte@127.0.0.1:5433/astrocyte'
+cd astrocyte-services-py/astrocyte-pgvector
+./scripts/migrate.sh
+```
+
+Requirements: **PostgreSQL 15+** (for `CREATE INDEX CONCURRENTLY IF NOT EXISTS`), **psql** on `PATH`.
+
+After migrations are applied, set **`bootstrap_schema: false`** in `vector_store_config` so the app does not run `CREATE TABLE` / indexes at runtime (see configuration table below). For a **single command** that starts Postgres, runs migrations, then starts the stack with runbook config, use **[`scripts/runbook-up.sh`](../scripts/runbook-up.sh)** (see **[Runbook](../README.md#runbook)**).
+
+**Embedding width:** [`migrations/002_astrocyte_vectors.sql`](migrations/002_astrocyte_vectors.sql) defines `vector(128)`. That must match **`embedding_dimensions`** in config. For another width, add a new migration (or edit before first deploy) and keep the Python config aligned.
+
+**Custom `table_name`:** The shipped SQL targets **`astrocyte_vectors`**. If you use another table name, copy and adjust the migration files accordingly.
+
+## Configuration
+
+| Constructor / YAML `vector_store_config` | Meaning |
+|--------------------------------------------|---------|
+| `dsn` | PostgreSQL connection URI (or set `DATABASE_URL` / `ASTROCYTES_PG_DSN`) |
+| `table_name` | Table name (default `astrocyte_vectors`; alphanumeric + underscore only) |
+| `embedding_dimensions` | Fixed `vector(N)` width; must match your embedding model and the **`vector(N)`** in SQL migrations (default **128**) |
+| `bootstrap_schema` | If **`true`** (default), create extension / table / btree index on first use (dev-friendly; no HNSW). If **`false`**, assume **`migrate.sh`** already applied [`migrations/`](migrations/) (production). |
+
+## How this fits `astrocyte_rest`
+
+1. **`astrocyte-py`** defines the **`VectorStore`** protocol and discovers adapters by **entry point** (`astrocyte.vector_stores`).
+2. **`astrocyte-pgvector`** registers **`pgvector` → `PgVectorStore`**. Installing this package makes the name **`pgvector`** available to **`resolve_provider()`**.
+3. **`astrocyte_rest/wiring.py`** calls **`resolve_vector_store(config)`**, which loads the class from the entry point and passes **`vector_store_config`** from YAML (or env-only defaults).
+4. **`astrocyte_rest/brain.py`** builds **`Astrocyte`** + **`PipelineOrchestrator`** with that store and your chosen **`llm_provider`** (still **`mock`** unless you configure a real LLM).
+
+Example **`ASTROCYTES_CONFIG_PATH`** snippet:
+
+```yaml
+provider_tier: storage
+vector_store: pgvector
+llm_provider: mock
+vector_store_config:
+  dsn: postgresql://astrocyte:astrocyte@127.0.0.1:5433/astrocyte
+  embedding_dimensions: 128
+  bootstrap_schema: false
+```
+
+Then run the REST service (from repo layout):
+
+```bash
+export ASTROCYTES_CONFIG_PATH=/path/to/that.yaml
+cd astrocyte-services-py/astrocyte-rest && uv run astrocyte-rest
+```
+
+Or set only env (no YAML file):
+
+```bash
+export ASTROCYTES_VECTOR_STORE=pgvector
+export DATABASE_URL=postgresql://astrocyte:astrocyte@127.0.0.1:5433/astrocyte
+# embedding_dimensions default 128 — override via YAML if you add a file
+cd astrocyte-services-py/astrocyte-rest && uv sync --extra pgvector
+```
+
+**Note:** `vector_store_config` for dimensions is only merged from YAML today; for env-only mode, add a small YAML or extend `brain.py` to pass `ASTROCYTES_EMBEDDING_DIMENSIONS` (future improvement).
+
+## Production notes
+
+- **HNSW** parameters (`m`, `ef_construction`) live in [`migrations/003_indexes.sql`](migrations/003_indexes.sql); tune with DBA guidance as load grows.
+- **Embedding dimension** must match the **`LLMProvider.embed()`** output used by the pipeline.
+- Use **secrets** for `dsn`, not committed YAML.

astrocyte_pgvector-0.1.0.dist-info/RECORD

@@ -0,0 +1,6 @@
+astrocyte_pgvector/__init__.py,sha256=RFaVN_DmwTwQg_dSlQor5JbT8GSDyuheJaN4zXKMLhc,139
+astrocyte_pgvector/store.py,sha256=qDd65lCsneDAAKy-iogZZVLM2R8UzOkdkeRT6vgUEgs,10874
+astrocyte_pgvector-0.1.0.dist-info/METADATA,sha256=wIqS_pfn-fBXWpEegOknKZwi60SboFdyIONBTaQMiSo,5342
+astrocyte_pgvector-0.1.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+astrocyte_pgvector-0.1.0.dist-info/entry_points.txt,sha256=1ddqMVkzitWuDUYGl_RtvQLeHPlagUY7yXpU0wL0AVU,76
+astrocyte_pgvector-0.1.0.dist-info/RECORD,,

astrocyte_pgvector-0.1.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

astrocyte_pgvector-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[astrocyte.vector_stores]
+pgvector = astrocyte_pgvector.store:PgVectorStore