Make pgvector and the model registry schema-aware; fix sqlite-vec's engine reconstruction

.github/workflows/ci.yml

@@ -6,21 +6,23 @@ on:
 jobs:
   label-gate:
     uses: AustralianCancerDataNetwork/cava-devops/.github/workflows/label-gate.yml@main
-  build-test:
-    uses: AustralianCancerDataNetwork/cava-devops/.github/workflows/build-test-postgres.yml@main
+  build-test-sqlite:
+    uses: AustralianCancerDataNetwork/cava-devops/.github/workflows/build-test.yml@main
+  build-test-postgres:
+    uses: AustralianCancerDataNetwork/cava-devops/.github/workflows/build-test-postgres-v2.yml@main
     with:
       postgres-image: pgvector/pgvector:pg16
       postgres-user: test
       postgres-password: test
       postgres-db: postgres
       setup-commands: |
         uv run omop-config configure omop_emb \
-          --set test_emb_db.kind=generic \
-          --set test_emb_db.connection.dialect=postgresql+psycopg \
-          --set test_emb_db.connection.host=localhost \
-          --set test_emb_db.connection.port=5432 \
-          --set test_emb_db.connection.user=test \
-          --set test_emb_db.connection.password=test \
-          --set test_emb_db.connection.database_name=test_omop_emb \
-          --set test_emb_db.connection.test_only=true \
-          --set test_emb_db.schema_name=public
+          --set test_emb_db_pg.kind=generic \
+          --set test_emb_db_pg.connection.dialect=postgresql+psycopg \
+          --set test_emb_db_pg.connection.host=localhost \
+          --set test_emb_db_pg.connection.port=5432 \
+          --set test_emb_db_pg.connection.user=test \
+          --set test_emb_db_pg.connection.password=test \
+          --set test_emb_db_pg.connection.database_name=test_omop_emb \
+          --set test_emb_db_pg.connection.test_only=true \
+          --set test_emb_db_pg.schema_name=public

docs/usage/backend-selection.md

@@ -10,7 +10,7 @@
 | **pgvector** | `pgvector` | `omop-emb[pgvector]` + PostgreSQL | Scales to large corpora. HNSW indexing and `halfvec` storage. |
 
 !!! note "FAISS is a sidecar, not a backend"
-    FAISS (`omop-emb[faiss-cpu]`) is a read-acceleration layer that sits on top of sqlite-vec or pgvector. It is not a primary backend and has no `backend_type` of its own. See the [CLI reference](cli.md#faiss-sidecar) for how to export and use FAISS indices.
+    FAISS (`omop-emb[faiss-cpu]`) is a read-acceleration layer that sits on top of sqlite-vec or pgvector. It is not a primary backend and has no `backend_type` of its own. See the [CLI reference](cli.md#build-faiss-cache) for how to export and use FAISS indices.
 
 ## Selecting a backend
 

docs/usage/installation.md

@@ -42,7 +42,7 @@ The backend (sqlite-vec or pgvector) is selected by a `[vector_stores.*]` entry'
 
 ```bash
 omop-config connections add emb --dialect postgresql+psycopg --host localhost --database-name omop_emb
-omop-config databases add emb_db --kind generic --connection emb
+omop-config databases add generic emb_db --connection emb
 omop-config vector-stores add vector_store --backend-type pgvector --database emb_db
 omop-config configure omop_emb --vector-store-name vector_store
 ```

docs/usage/interface-guide.md

@@ -28,14 +28,19 @@ backend = resolve_backend_from_resolved_vector_store(resolved)
 Or construct one directly:
 
 ```python
-from omop_emb.backends.sqlitevec import SQLiteVecEmbeddingBackend
+from sqlalchemy import create_engine
+from omop_emb.backends.sqlitevec import SQLiteVecEmbeddingBackend, create_sqlitevec_engine
 from omop_emb.backends.pgvector import PGVectorEmbeddingBackend
 
 # sqlite-vec
-backend = SQLiteVecEmbeddingBackend.from_path(db_path="/data/omop_emb.db")
+backend = SQLiteVecEmbeddingBackend(
+    emb_engine=create_sqlitevec_engine(create_engine("sqlite:///data/omop_emb.db"))
+)
 
 # pgvector
-backend = PGVectorEmbeddingBackend.from_db_url(db_url="postgresql+psycopg://user:pass@host:5432/db")
+backend = PGVectorEmbeddingBackend(
+    emb_engine=create_engine("postgresql+psycopg://user:pass@host:5432/db")
+)
 ```
 
 ---
@@ -186,9 +191,14 @@ results = reader.get_nearest_concepts(query_embedding=joint_vec[None, :], k=10)
 `get_nearest_concepts_from_query_texts` takes a `ModelBackend` directly: build one with `omop_llm.build_model_backend` (the reader has no default backend of its own to embed with):
 
 ```python
-from omop_llm import build_model_backend
+from omop_llm import Capabilities, build_model_backend
 
-model_backend = build_model_backend("ollama", "nomic-embed-text:v1.5", base_url="http://localhost:11434")
+model_backend = build_model_backend(
+    "ollama",
+    "nomic-embed-text:v1.5",
+    model_capabilities=Capabilities(embeddings=True),
+    base_url="http://localhost:11434",
+)
 
 results = reader.get_nearest_concepts_from_query_texts(
     query_texts=("high blood pressure", "type 2 diabetes"),
@@ -267,6 +277,7 @@ from omop_llm import build_model_backend
 model_backend = build_model_backend(
     "ollama",
     "nomic-embed-text:v1.5",
+    model_capabilities=Capabilities(embeddings=True),
     base_url="http://host.docker.internal:11434",
 )
 
@@ -279,6 +290,7 @@ print(model_backend.dimensions())  # auto-discovered via Ollama /api/show
 model_backend = build_model_backend(
     "openai",
     "text-embedding-3-large",
+    model_capabilities=Capabilities(embeddings=True),
     base_url="https://api.openai.com/v1",
     api_key="sk-...",
 )

pyproject.toml

@@ -102,7 +102,7 @@ markers = [
   "faiss: FAISS sidecar tests",
 ]
 testpaths = ["tests"]
-addopts = "-v --tb=short"
+addopts = "-v --tb=short -m 'not db_dialect'"
 filterwarnings = [
     "ignore:datetime\\.datetime\\.utcfromtimestamp\\(\\) is deprecated:DeprecationWarning:dateutil.tz.tz",
     "ignore:builtin type .* has no __module__ attribute:DeprecationWarning",

src/omop_emb/backends/base_backend.py

@@ -6,7 +6,12 @@
 from datetime import datetime
 from typing import Any, Callable, Generic, Iterable, Mapping, Optional, Sequence, Tuple, TypeVar, Union
 from numpy import ndarray
-from oa_configurator import ResolvedDatabase, ResolvedVectorStore
+from oa_configurator import (
+    SCHEMA_TRANSLATE_MAP_KEY,
+    Dialect,
+    ResolvedDatabase,
+    ResolvedVectorStore,
+)
 from sqlalchemy import Engine
 from sqlalchemy.engine import make_url
 from sqlalchemy.orm import sessionmaker
@@ -22,7 +27,12 @@
 
 from omop_emb.backends.embedding_table import ConceptEmbeddingRecord
 from omop_emb.backends.index_config import IndexConfig, FlatIndexConfig
-from omop_emb.model_registry import EmbeddingModelRecord, RegistryManager
+from omop_emb.model_registry import (
+    EmbeddingModelRecord,
+    REGISTRY_SCHEMA_KEY,
+    RegistryManager,
+    resolve_registry_physical_schema,
+)
 from omop_emb.utils.embedding_utils import (
     EmbeddingConceptFilter,
     NearestConceptMatch,
@@ -168,15 +178,21 @@ class EmbeddingBackend(ABC, Generic[TEmbeddingTable]):
 
     DEFAULT_K_NEAREST = 10
 
-    def __init__(self, emb_engine: Engine) -> None:
+    def __init__(
+        self,
+        emb_engine: Engine,
+        *,
+        resolved: ResolvedDatabase | None = None,
+    ) -> None:
         actual_dialect = emb_engine.dialect.name
         if actual_dialect != self.dialect:
             raise ValueError(
                 f"{type(self).__name__} requires a '{self.dialect}'-dialect engine, "
                 f"got '{actual_dialect}'."
             )
         super().__init__()
-        self._registry = RegistryManager(emb_engine)
+        self._resolved = resolved
+        self._registry = RegistryManager(emb_engine, resolved=resolved)
         self._table_cache: dict[str, TEmbeddingTable] = {}
         self._initialise_store()
 
@@ -278,7 +294,9 @@ def register_model(
         ModelRegistrationConflictError
             If the model is already registered with a different dimensionality.
         ValueError
-            If ``metadata`` contains a reserved key.
+            If ``metadata`` contains a reserved key, or if ``index_config`` is
+            not ``FlatIndexConfig()`` (non-FLAT indexes may only be built
+            after registration, not at registration time).
         """
 
         if index_config is None:
@@ -1036,20 +1054,29 @@ def resolve_backend(
 
     dialect = make_url(database.connection.url).get_backend_name()
 
+    # The model registry lives in its own reserved schema (MODEL_REGISTRY_SCHEMA),
+    # independent of database's own schema. create_engine() merges this extra
+    # key onto its own configured map rather than replacing it.
+    registry_schema = resolve_registry_physical_schema(database.connection.dialect_name)
+    registry_schema_translate_map = {REGISTRY_SCHEMA_KEY: registry_schema}
+
     if resolved_backend == BackendType.SQLITEVEC:
-        from omop_emb.backends.sqlitevec import SQLiteVecEmbeddingBackend
+        from omop_emb.backends.sqlitevec import SQLiteVecEmbeddingBackend, create_sqlitevec_engine
 
-        if dialect != "sqlite":
+        if dialect != Dialect.SQLITE:
             raise RuntimeError(
                 f"sqlitevec backend requires a sqlite-dialect database, got dialect: {dialect!r}."
             )
-        db_path = make_url(database.connection.url).database
-        assert db_path is not None, "ConnectionConfig.build_url() always sets a database segment for sqlite"
-        logger.info(f"Using SQLiteVec backend with database file: {db_path}")
-        return SQLiteVecEmbeddingBackend.from_path(db_path)
+        emb_engine = create_sqlitevec_engine(
+            database.create_engine(
+                execution_options={SCHEMA_TRANSLATE_MAP_KEY: registry_schema_translate_map}
+            )
+        )
+        logger.info(f"Using SQLiteVec backend with engine: {emb_engine.url}")
+        return SQLiteVecEmbeddingBackend(emb_engine=emb_engine, resolved=database)
 
     if resolved_backend == BackendType.PGVECTOR:
-        if dialect != "postgresql":
+        if dialect != Dialect.POSTGRESQL:
             raise RuntimeError(
                 "The resolved URL must point to a PostgreSQL database "
                 f"(pgvector extension required), got dialect: {dialect!r}."
@@ -1061,9 +1088,11 @@ def resolve_backend(
                 "pgvector backend is not installed. "
                 "Install it with: pip install omop-emb[pgvector]"
             ) from exc
-        emb_engine = database.create_engine()
+        emb_engine = database.create_engine(
+            execution_options={SCHEMA_TRANSLATE_MAP_KEY: registry_schema_translate_map}
+        )
         logger.info(f"Using pgvector backend with engine: {emb_engine.url}")
-        return PGVectorEmbeddingBackend(emb_engine=emb_engine)
+        return PGVectorEmbeddingBackend(emb_engine=emb_engine, resolved=database)
 
     raise RuntimeError(f"Implementation for {resolved_backend.value} is not available.")
 

src/omop_emb/backends/db_utils.py

@@ -3,6 +3,7 @@
 from contextlib import contextmanager
 from typing import Any, Iterator, Sequence
 
+from oa_configurator import Dialect
 from sqlalchemy import Column, Select, text
 from sqlalchemy.orm import Session
 from sqlalchemy.sql.base import ColumnCollection
@@ -100,13 +101,13 @@ def temp_filter_table(
     connections safely.  The table is truncated before each use and cleaned up
     when the connection is returned to the pool.
     """
-    if dialect == "postgresql":
+    if dialect == Dialect.POSTGRESQL:
         session.execute(
             text(
                 f'CREATE TEMPORARY TABLE "{table_name}" (id {col_type}) ON COMMIT DROP'
             )
         )
-    elif dialect == "sqlite":
+    elif dialect == Dialect.SQLITE:
         session.execute(
             text(f'CREATE TEMPORARY TABLE IF NOT EXISTS "{table_name}" (id {col_type})')
         )

src/omop_emb/backends/index_config.py

@@ -126,9 +126,7 @@ def from_dict(cls, config_dict: dict[str, Any]) -> Self:
         Notes
         -----
         Use this method to reconstruct an ``IndexConfig`` from the ORM
-        ``index_config`` JSON column. It is distinct from
-        :meth:`from_metadata`, which reads from a metadata dict that wraps the
-        config under ``"index_config"`` key.
+        ``index_config`` JSON column.
         """
         if not is_dataclass(cls):
             raise TypeError(f"Must be called on a dataclass, not {cls.__name__}.")

src/omop_emb/backends/pgvector/pg_backend.py

@@ -5,7 +5,8 @@
 from typing import Mapping, Optional, Sequence, Tuple
 
 from numpy import ndarray
-from sqlalchemy import Engine, select, text, create_engine
+from oa_configurator import Dialect, ResolvedDatabase
+from sqlalchemy import Engine, select, text
 
 try:
     from pgvector.sqlalchemy import Vector  # noqa: F401
@@ -78,25 +79,14 @@ class PGVectorEmbeddingBackend(EmbeddingBackend[type[PGEmbeddingTable]]):
     ``__init__``.
     """
 
-    def __init__(self, emb_engine: Engine) -> None:
+    def __init__(
+        self,
+        emb_engine: Engine,
+        *,
+        resolved: ResolvedDatabase | None = None,
+    ) -> None:
         self._index_managers: dict[str, PGVectorBaseIndexManager] = {}
-        super().__init__(emb_engine=emb_engine)
-
-    @classmethod
-    def from_db_url(cls, db_url: str) -> PGVectorEmbeddingBackend:
-        """Create a pgvector embedding backend from a database URL.
-
-        Parameters
-        ----------
-        db_url : str
-            Database URL in SQLAlchemy format, e.g. ``postgresql://user:pass@host:port/dbname``.
-
-        Returns
-        -------
-        PGVectorEmbeddingBackend
-        """
-        engine = create_engine(db_url, echo=False)
-        return cls(emb_engine=engine)
+        super().__init__(emb_engine=emb_engine, resolved=resolved)
 
     # ------------------------------------------------------------------
     # Backend identity
@@ -108,7 +98,7 @@ def backend_type(self) -> BackendType:
 
     @property
     def dialect(self) -> str:
-        return "postgresql"
+        return Dialect.POSTGRESQL
 
     # ------------------------------------------------------------------
     # Store lifecycle
@@ -135,7 +125,9 @@ def _create_storage_table(
         self, model_record: EmbeddingModelRecord
     ) -> type[PGEmbeddingTable]:
         return create_pg_embedding_table(
-            engine=self.emb_engine, model_record=model_record
+            engine=self.emb_engine,
+            model_record=model_record,
+            resolved=self._resolved,
         )
 
     def _delete_storage_table(self, model_record: EmbeddingModelRecord) -> None:
@@ -181,7 +173,11 @@ def register_model(
         Raises
         ------
         ValueError
-            If ``dimensions`` exceeds the pgvector halfvec limit of 4 000.
+            If ``dimensions`` exceeds the pgvector halfvec limit of 4 000, or
+            for any reason :meth:`EmbeddingBackend.register_model` raises
+            (a reserved metadata key, or a non-``FlatIndexConfig`` index).
+        ModelRegistrationConflictError
+            If the model is already registered with a different dimensionality.
         """
         vector_column_type_for_dimensions(dimensions)  # validates halfvec limit
         return super().register_model(

src/omop_emb/backends/pgvector/pg_index_manager.py

@@ -17,6 +17,7 @@
 import logging
 from typing import Generic, TypeVar
 
+from oa_configurator import qualified, physical_schema_of
 from sqlalchemy import Engine, inspect, text
 
 from omop_emb.config import IndexType, MetricType, VectorColumnType
@@ -77,7 +78,10 @@ def index_config(self) -> C:
     def has_index(self, metric_type: MetricType) -> bool:
         with self._engine.connect() as conn:
             existing = {
-                idx["name"] for idx in inspect(conn).get_indexes(self._tablename)
+                idx["name"]
+                for idx in inspect(conn).get_indexes(
+                    self._tablename, schema=physical_schema_of(conn)
+                )
             }
         return self._index_name(metric_type) in existing
 
@@ -101,7 +105,9 @@ def drop_index(self, metric_type: MetricType) -> None:
         name = self._index_name(metric_type)
         existed = self.has_index(metric_type)
         with self._engine.begin() as conn:
-            conn.execute(text(f"DROP INDEX IF EXISTS {name}"))
+            conn.execute(
+                text(f"DROP INDEX IF EXISTS {qualified(conn, name, physical_schema=physical_schema_of(conn))}")
+            )
         if existed:
             logger.info(f"Dropped pgvector index '{name}'.")
 
@@ -190,9 +196,10 @@ def supported_index_type(self) -> IndexType:
     def _create_index_ddl(self, metric_type: MetricType) -> str:
         ops = self._ops_for_metric(metric_type)
         cfg = self.index_config
+        table_ref = qualified(self._engine, self._tablename, physical_schema=physical_schema_of(self._engine))
         return (
             f"CREATE INDEX {self._index_name(metric_type)} "
-            f"ON {self._tablename} "
+            f"ON {table_ref} "
             f"USING hnsw ({self._embedding_column} {ops}) "
             f"WITH (m = {cfg.num_neighbors}, ef_construction = {cfg.ef_construction})"
         )