fix: use VECTOR_STORE.DIMENSIONS for SQLAlchemy vector column size#465
fix: use VECTOR_STORE.DIMENSIONS for SQLAlchemy vector column size#465paralizeer wants to merge 4 commits intoplastic-labs:mainfrom
Conversation
The Document and MessageEmbedding models hardcoded Vector(1536) which
assumes OpenAI text-embedding-3-small dimensions. This causes runtime
errors ('expected 1536 dimensions, not 768') when using embedding
providers that produce different dimensionalities (e.g. nomic-embed-text
produces 768 dims, many local models use 384 or 1024).
Fix: replace hardcoded Vector(1536) with Vector(settings.VECTOR_STORE.DIMENSIONS)
so the SQLAlchemy column type matches the configured embedding dimensions.
The VECTOR_STORE.DIMENSIONS setting already exists (default: 1536) and
is used for LanceDB/TurboPuffer — this extends it to the pgvector models.
For custom embedding models, set in config.toml:
[vector_store]
DIMENSIONS = 768 # match your embedding model's output dimensions
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughReplaced hard-coded 1536 vector dimensionality with Changes
Sequence Diagram(s)sequenceDiagram
participant App as App (lifespan)
participant Cache as Cache init
participant DB as Postgres DB
participant Check as check_vector_dimensions
participant Logger as Logger
App->>Cache: initialize cache
Cache-->>App: success/failure handling
App->>DB: open async session (SessionLocal)
DB-->>App: session
App->>Check: await check_vector_dimensions(session)
Check->>DB: query pg_catalog (pg_attribute/pg_class/pg_namespace)
DB-->>Check: atttypmod rows for embedding columns
alt settings.VECTOR_STORE.TYPE == "pgvector"
Check->>Check: extract stored_dim from atttypmod
alt stored_dim == settings.VECTOR_STORE.DIMENSIONS
Check->>Logger: log OK (dimensions match)
Check-->>App: return None
else stored_dim != settings.VECTOR_STORE.DIMENSIONS
Check->>Logger: log mismatch
Check-->>App: raise RuntimeError
end
else
Check-->>App: return immediately (non-pgvector backend)
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://cdn.statically.io/img/avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/models.py`:
- Line 390: The embedding dimensionality is hardcoded to 1536 in
embedding_client.py which can mismatch the configured VECTOR_STORE.DIMENSIONS
used by the ORM column definition (the embedding field: "embedding:
MappedColumn[Any] = mapped_column(Vector(settings.VECTOR_STORE.DIMENSIONS),
nullable=True)"). Update embedding_client.py to replace all occurrences of the
literal 1536 (e.g., the "output_dimensionality" values) with
settings.VECTOR_STORE.DIMENSIONS so Gemini embeddings respect the configured
vector size; also break the long models.py embedding column line into two
shorter parts to keep it under 88 characters (split the
mapped_column(Vector(...), nullable=True) expression across lines while
preserving the same expression and indentation).
- Line 282: The embedding column declaration uses
settings.VECTOR_STORE.DIMENSIONS but migrations hardcode Vector(1536), and the
long line exceeds 88 chars; update the model's embedding declaration formatting
so it wraps (use mapped_column(Vector(settings.VECTOR_STORE.DIMENSIONS),
nullable=True) across multiple lines) and add a startup validation function
(e.g., check_vector_dimension or in app initialization) that queries the DB
column type for the embeddings column and compares its vector dimension to
settings.VECTOR_STORE.DIMENSIONS, failing fast with a clear error if they differ
(or log and abort), so runtime configuration and the existing migrations cannot
diverge silently.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
- Replace all output_dimensionality: 1536 in embedding_client.py with settings.VECTOR_STORE.DIMENSIONS so Gemini embeddings respect config - Wrap long embedding column declarations for readability (< 88 chars) - Add check_vector_dimensions() startup validator that queries DB column type and fails fast if Alembic migrations diverge from runtime config Co-authored-by: CodeRabbit <coderabbit@coderabbit.ai>
|
Applied both suggestions from @coderabbitai:
The validator is async and takes a session, so it can be called from the app lifespan on first DB connection. Happy to wire it into the startup sequence if the team prefers that over leaving it as a utility. |
|
Tip For best results, initiate chat on the files or code changes. 🧩 Analysis chain🏁 Script executed: Length of output: 4639 🏁 Script executed: Length of output: 25136 🏁 Script executed: Length of output: 4974
1. Two gaps in
|
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/embedding_client.py (1)
90-94:⚠️ Potential issue | 🟠 MajorAdd
dimensionsparameter to OpenAI embeddings.create() callOpenAI's
text-embedding-3-smallsupports a configurabledimensionsparameter. The current code doesn't pass it, causing a mismatch whenVECTOR_STORE.DIMENSIONSis configured to a non-default value (e.g., 768). Since the database schema enforces the configured dimensions and Gemini already respects this setting viaoutput_dimensionality, OpenAI and OpenRouter should also passdimensions=settings.VECTOR_STORE.DIMENSIONStoembeddings.create()for consistency.Current code (lines 90-94)
else: # openai response = await self.client.embeddings.create( model=self.model, input=query ) return response.data[0].embedding🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/embedding_client.py` around lines 90 - 94, The OpenAI embeddings call in the else branch of the embedding creation doesn't pass the configured vector dimensionality; update the call to self.client.embeddings.create to include dimensions=settings.VECTOR_STORE.DIMENSIONS (i.e., await self.client.embeddings.create(model=self.model, input=query, dimensions=settings.VECTOR_STORE.DIMENSIONS)) and ensure settings is imported/available in the scope where this method/class (the embedding client method using self.client.embeddings.create, self.model, and query) is defined so the returned embedding matches the DB schema.
🧹 Nitpick comments (4)
src/models.py (2)
590-590: Remove shadowed logger declaration.Line 35 already declares a module-level
logger = getLogger(__name__). This local re-declaration shadows it unnecessarily.♻️ Proposed fix
async def check_vector_dimensions(session: Any) -> None: """ Validate that the DB vector column dimension matches VECTOR_STORE.DIMENSIONS. Call this at application startup to fail fast if Alembic migrations and runtime config have diverged. """ - logger = getLogger(__name__) expected = settings.VECTOR_STORE.DIMENSIONS🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/models.py` at line 590, Remove the redundant local declaration `logger = getLogger(__name__)` that shadows the existing module-level `logger`; instead, delete that local line and use the already-declared module-level `logger` everywhere in the surrounding function/class (keep existing log calls unchanged). Ensure no other local variable named `logger` remains in the scope so all logging uses the single module-level `logger`.
603-603: Line exceeds 88 characters.As per coding guidelines, use line length of 88 characters (Black compatible).
♻️ Proposed fix
if row is None: - logger.warning("check_vector_dimensions: could not find message_embeddings.embedding column") + logger.warning( + "check_vector_dimensions: could not find message_embeddings.embedding column" + ) return🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/models.py` at line 603, The warning message string in logger.warning call inside check_vector_dimensions is longer than 88 characters; shorten or rephrase the message (e.g., "check_vector_dimensions: missing message_embeddings.embedding column") or split into a shorter message so the logger.warning("...") call in that function stays within the 88-character limit.src/embedding_client.py (2)
255-255: Same line length issue as above.♻️ Proposed fix
if isinstance(self.client, genai.Client): response = await self.client.aio.models.embed_content( model=self.model, contents=[item.text for item in batch], - config={"output_dimensionality": settings.VECTOR_STORE.DIMENSIONS}, + config={ + "output_dimensionality": settings.VECTOR_STORE.DIMENSIONS + }, )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/embedding_client.py` at line 255, The line setting config={"output_dimensionality": settings.VECTOR_STORE.DIMENSIONS} is too long; break the dictionary across multiple lines to respect max line length and match surrounding formatting. Update the call where this appears (the argument named config in the embedding client call) to use a multi-line dict with the key on its own line and the value on the next (or align with other kwargs), and run the formatter/linter so the change is consistent with the project's style.
119-119: Line may exceed 88 characters due to indentation depth.Consider breaking the config argument across lines for Black compliance.
♻️ Proposed fix
if isinstance(self.client, genai.Client): # Type cast needed due to genai type signature complexity response = await self.client.aio.models.embed_content( model=self.model, contents=batch, # pyright: ignore[reportArgumentType] - config={"output_dimensionality": settings.VECTOR_STORE.DIMENSIONS}, + config={ + "output_dimensionality": settings.VECTOR_STORE.DIMENSIONS + }, )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/embedding_client.py` at line 119, The line with the config argument (config={"output_dimensionality": settings.VECTOR_STORE.DIMENSIONS}) is too long; split the dict across multiple lines so Black will wrap it. Replace the single-line dict with a multi-line form such as config={ "output_dimensionality": settings.VECTOR_STORE.DIMENSIONS, } or break the call so the config= argument is on its own line and the dict entries are indented on subsequent lines; update the call site where config is passed to preserve trailing commas and indentation.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/models.py`:
- Around line 584-612: check_vector_dimensions only queries the
message_embeddings.embedding column; extend it to validate the
documents.embedding column too (both tables use
settings.VECTOR_STORE.DIMENSIONS) by iterating over the table names (e.g.,
["message_embeddings", "documents"]) and running the same pg_attribute query per
table, comparing each atttypmod to settings.VECTOR_STORE.DIMENSIONS and raising
the same RuntimeError if any mismatch is found; keep existing logging behavior
(warning if column not found, info when OK) and reference the same column name
"embedding" and constant settings.VECTOR_STORE.DIMENSIONS in the updated
check_vector_dimensions implementation.
---
Outside diff comments:
In `@src/embedding_client.py`:
- Around line 90-94: The OpenAI embeddings call in the else branch of the
embedding creation doesn't pass the configured vector dimensionality; update the
call to self.client.embeddings.create to include
dimensions=settings.VECTOR_STORE.DIMENSIONS (i.e., await
self.client.embeddings.create(model=self.model, input=query,
dimensions=settings.VECTOR_STORE.DIMENSIONS)) and ensure settings is
imported/available in the scope where this method/class (the embedding client
method using self.client.embeddings.create, self.model, and query) is defined so
the returned embedding matches the DB schema.
---
Nitpick comments:
In `@src/embedding_client.py`:
- Line 255: The line setting config={"output_dimensionality":
settings.VECTOR_STORE.DIMENSIONS} is too long; break the dictionary across
multiple lines to respect max line length and match surrounding formatting.
Update the call where this appears (the argument named config in the embedding
client call) to use a multi-line dict with the key on its own line and the value
on the next (or align with other kwargs), and run the formatter/linter so the
change is consistent with the project's style.
- Line 119: The line with the config argument (config={"output_dimensionality":
settings.VECTOR_STORE.DIMENSIONS}) is too long; split the dict across multiple
lines so Black will wrap it. Replace the single-line dict with a multi-line form
such as config={ "output_dimensionality": settings.VECTOR_STORE.DIMENSIONS, } or
break the call so the config= argument is on its own line and the dict entries
are indented on subsequent lines; update the call site where config is passed to
preserve trailing commas and indentation.
In `@src/models.py`:
- Line 590: Remove the redundant local declaration `logger =
getLogger(__name__)` that shadows the existing module-level `logger`; instead,
delete that local line and use the already-declared module-level `logger`
everywhere in the surrounding function/class (keep existing log calls
unchanged). Ensure no other local variable named `logger` remains in the scope
so all logging uses the single module-level `logger`.
- Line 603: The warning message string in logger.warning call inside
check_vector_dimensions is longer than 88 characters; shorten or rephrase the
message (e.g., "check_vector_dimensions: missing message_embeddings.embedding
column") or split into a shorter message so the logger.warning("...") call in
that function stays within the 88-character limit.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: a44177b4-5832-4208-97a3-73adc8e522e5
📒 Files selected for processing (2)
src/embedding_client.pysrc/models.py
- Check both message_embeddings.embedding and documents.embedding - Gate on VECTOR_STORE.TYPE == 'pgvector' (skip for LanceDB/TurboPuffer) - Use parameterized queries instead of f-strings - Wire into app lifespan in main.py via SessionLocal so the check runs at startup before serving requests Co-authored-by: CodeRabbit <coderabbit@coderabbit.ai>
|
Updated with both fixes:
|
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/main.py`:
- Around line 136-141: The vector-dimension check is currently executed before
entering the lifespan try/finally, so if check_vector_dimensions(session) raises
the cleanup in the finally block never runs; move the async with SessionLocal()
and await check_vector_dimensions(session) inside the lifespan try block (i.e.,
into the same try that surrounds yield) so any exception during
check_vector_dimensions still triggers the finally cleanup — adjust placement
around the existing try/yield/finally and keep using SessionLocal and
check_vector_dimensions as the identifying symbols.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 886a45bf-142c-44ee-8abf-3045e24384f8
📒 Files selected for processing (2)
src/main.pysrc/models.py
🚧 Files skipped from review as they are similar to previous changes (1)
- src/models.py
Ensures cleanup (vector store, cache, engine disposal) always runs even if the dimension check raises at startup. Co-authored-by: CodeRabbit <coderabbit@coderabbit.ai>
|
Fixed — moved Also note: the |
1 participant
Problem
DocumentandMessageEmbeddingmodels hardcodeVector(1536):This assumes OpenAI
text-embedding-3-smalldimensions. Using any other embedding model produces a runtime error:Fix
Replace
Vector(1536)withVector(settings.VECTOR_STORE.DIMENSIONS).VECTOR_STORE.DIMENSIONSalready exists (default: 1536) and is used for LanceDB/TurboPuffer — this extends it consistently to the pgvector SQLAlchemy models.Impact
No breaking change for default deployments (DIMENSIONS defaults to 1536). Users with non-OpenAI embedding models can now set:
Summary by CodeRabbit