From 01dfec523d2e9737dfac6366dac44ea2d126505c Mon Sep 17 00:00:00 2001
From: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>
Date: Tue, 12 May 2026 17:23:57 -0400
Subject: [PATCH] feat(models): honor EMBEDDING_VECTOR_DIMENSIONS in pgvector
 columns

---
 src/models.py                   |  7 +++-
 tests/test_models_vector_dim.py | 69 +++++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+), 2 deletions(-)
 create mode 100644 tests/test_models_vector_dim.py

diff --git a/src/models.py b/src/models.py
index c2aa84dd..6433225a 100644
--- a/src/models.py
+++ b/src/models.py
@@ -25,12 +25,15 @@ from sqlalchemy.orm import Mapped, MappedColumn, mapped_column, relationship
 from sqlalchemy.sql import func
 from typing_extensions import override
 
+from src.config import settings
 from src.utils.types import DocumentLevel, TaskType, VectorSyncState
 
 from .db import Base
 
 load_dotenv(override=True)
 
+_VECTOR_DIM: int = settings.EMBEDDING.VECTOR_DIMENSIONS
+
 logger = getLogger(__name__)
 
 
@@ -278,7 +281,7 @@ class MessageEmbedding(Base):
         BigInteger, Identity(), primary_key=True, autoincrement=True
     )
     content: Mapped[str] = mapped_column(TEXT)
-    embedding: MappedColumn[Any] = mapped_column(Vector(1536), nullable=True)
+    embedding: MappedColumn[Any] = mapped_column(Vector(_VECTOR_DIM), nullable=True)
     message_id: Mapped[str] = mapped_column(
         ForeignKey("messages.public_id", ondelete="CASCADE"), nullable=False, index=True
     )
@@ -386,7 +389,7 @@ class Document(Base):
     times_derived: Mapped[int] = mapped_column(
         Integer, nullable=False, server_default=text("1")
     )
-    embedding: MappedColumn[Any] = mapped_column(Vector(1536), nullable=True)
+    embedding: MappedColumn[Any] = mapped_column(Vector(_VECTOR_DIM), nullable=True)
     source_ids: Mapped[list[str] | None] = mapped_column(
         JSONB, nullable=True, server_default=text("NULL")
     )
diff --git a/tests/test_models_vector_dim.py b/tests/test_models_vector_dim.py
new file mode 100644
index 00000000..2233b840
--- /dev/null
+++ b/tests/test_models_vector_dim.py
@@ -0,0 +1,69 @@
+"""Phase 1: verify src/models.py honors EMBEDDING_VECTOR_DIMENSIONS at import time."""
+
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+_PROJECT_ROOT = Path(__file__).resolve().parent.parent
+
+
+def _run_in_fresh_interpreter(env_overrides: dict[str, str]) -> dict[str, int]:
+    """Import src.models in a fresh interpreter and return the vector dims.
+
+    A subprocess is required because src.models reads
+    settings.EMBEDDING.VECTOR_DIMENSIONS at module import time to construct
+    SQLAlchemy column types — reloading the module in-process would conflict
+    with the existing Base.registry from earlier imports.
+
+    PYTHON_DOTENV_DISABLED=1 prevents config.py:20 from reloading the
+    developer's .env file (which calls load_dotenv with override=True)
+    and clobbering our test overrides.
+    """
+    env: dict[str, str] = {
+        **os.environ,
+        "PYTHON_DOTENV_DISABLED": "1",
+        **env_overrides,
+    }
+    snippet = (
+        "import json\n"
+        "from src.models import Document, MessageEmbedding\n"
+        "print(json.dumps({\n"
+        "    'message_embedding_dim': MessageEmbedding.__table__.c.embedding.type.dim,\n"
+        "    'document_dim': Document.__table__.c.embedding.type.dim,\n"
+        "}))\n"
+    )
+    result = subprocess.run(
+        [sys.executable, "-c", snippet],
+        env=env,
+        cwd=str(_PROJECT_ROOT),
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+    parsed: dict[str, int] = json.loads(result.stdout.strip().splitlines()[-1])
+    return parsed
+
+
+def test_models_uses_default_1536_when_no_env_override() -> None:
+    dims = _run_in_fresh_interpreter({})
+    assert dims == {"message_embedding_dim": 1536, "document_dim": 1536}
+
+
+def test_models_honors_explicit_embedding_vector_dimensions() -> None:
+    # Phase 1 still respects the dim-vs-MIGRATED guard at src/config.py:1278.
+    # Set VECTOR_STORE_TYPE=lancedb + VECTOR_STORE_MIGRATED=true to satisfy it.
+    # lancedb is chosen over turbopuffer because turbopuffer requires an
+    # additional VECTOR_STORE_TURBOPUFFER_API_KEY env var. Phase 2 deletes
+    # the guard and these escape-hatch envs become unnecessary.
+    dims = _run_in_fresh_interpreter(
+        {
+            "EMBEDDING_VECTOR_DIMENSIONS": "768",
+            "VECTOR_STORE_TYPE": "lancedb",
+            "VECTOR_STORE_MIGRATED": "true",
+        }
+    )
+    assert dims == {"message_embedding_dim": 768, "document_dim": 768}