Skip to content

fix(migrations): dynamic embedding dimenstions #770

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
MariusMueller wants to merge 2 commits into
base: main
Choose a base branch
from
Open

fix(migrations): dynamic embedding dimenstions #770

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
78de142
fix(migrations): derive embedding dimenstions from settings instead o…
MariusMueller Jun 3, 2026
5d31cb7
fix(llm): add model "mistral-embed" to known models rejecting dimenst…
MariusMueller Jun 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions migrations/versions/119a52b73c60_support_external_embeddings.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
from pgvector.sqlalchemy import Vector

from migrations.utils import column_exists, constraint_exists, get_schema, index_exists
from src.config import settings

# revision identifiers, used by Alembic.
revision: str = "119a52b73c60"
Expand All @@ -42,15 +43,15 @@ def upgrade() -> None:
op.alter_column(
"message_embeddings",
"embedding",
existing_type=Vector(1536),
existing_type=Vector(settings.EMBEDDING.VECTOR_DIMENSIONS),
nullable=True,
schema=schema,
)

op.alter_column(
"documents",
"embedding",
existing_type=Vector(1536),
existing_type=Vector(settings.EMBEDDING.VECTOR_DIMENSIONS),
nullable=True,
schema=schema,
)
Comment on lines 43 to 57

@coderabbitai coderabbitai Bot Jun 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

existing_type mismatch risk if config differs from deployed schema.

The existing_type parameter tells Alembic the current column type. If settings.EMBEDDING.VECTOR_DIMENSIONS doesn't match the dimension actually in the database (e.g., the initial migration ran with a different config value), this could cause:

  • Incorrect DDL generation on databases that use existing_type for ALTER statements
  • Silent type coercion issues

Since this migration only changes nullable, PostgreSQL typically handles this correctly regardless of existing_type, but the mismatch creates fragility. Consider hardcoding the expected dimension or validating it matches the database at runtime.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@migrations/versions/119a52b73c60_support_external_embeddings.py` around lines
43 - 57, The migration uses op.alter_column on "message_embeddings" and
"documents" with existing_type=Vector(settings.EMBEDDING.VECTOR_DIMENSIONS),
which risks mismatch if the deployed DB vector dimension differs from current
config; change the migration to either hardcode the expected vector dimension
(use the literal integer previously used for this schema) in
existing_type=Vector(<N>) for both op.alter_column calls or add a runtime
validation step that queries the DB column type/dimension and raises/adjusts if
it differs before calling op.alter_column; refer to the existing symbols
op.alter_column, existing_type, Vector, settings.EMBEDDING.VECTOR_DIMENSIONS and
the two table names to locate where to apply the fix.

Expand Down
2 changes: 1 addition & 1 deletion migrations/versions/917195d9b5e9_add_messageembedding_table.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ def upgrade() -> None:
"message_embeddings",
sa.Column("id", sa.BigInteger(), sa.Identity(), nullable=False),
sa.Column("content", sa.Text(), nullable=False),
sa.Column("embedding", Vector(1536), nullable=False),
sa.Column("embedding", Vector(settings.EMBEDDING.VECTOR_DIMENSIONS), nullable=False),
sa.Column("message_id", sa.Text(), nullable=False),
sa.Column("workspace_name", sa.Text(), nullable=False),
sa.Column("session_name", sa.Text(), nullable=True),
Expand Down
2 changes: 1 addition & 1 deletion migrations/versions/a1b2c3d4e5f6_initial_schema.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -363,7 +363,7 @@ def upgrade() -> None:
server_default="{}",
),
sa.Column("content", sa.Text(), nullable=False),
sa.Column("embedding", Vector(1536), nullable=True), # pyright: ignore
sa.Column("embedding", Vector(settings.EMBEDDING.VECTOR_DIMENSIONS), nullable=True), # pyright: ignore

@coderabbitai coderabbitai Bot Jun 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | 🏗️ Heavy lift

Configuration-driven migrations break reproducibility.

Alembic migrations should be immutable historical records. Reading settings.EMBEDDING.VECTOR_DIMENSIONS at runtime means:

  1. The same migration produces different schemas in different environments if EMBEDDING_VECTOR_DIMENSIONS varies
  2. Re-running migrations after a config change could fail (dimension mismatch with existing data) or create inconsistent state
  3. Migration history no longer documents what schema was actually applied

If you need to support multiple vector dimensions, consider:

  • Keeping hardcoded values in migrations for determinism (document the expected dimension in migration docstring)
  • Creating separate migration branches per deployment scenario
  • Using a dedicated "resize embeddings" migration when dimensions change
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@migrations/versions/a1b2c3d4e5f6_initial_schema.py` at line 366, The
migration is reading settings.EMBEDDING.VECTOR_DIMENSIONS at runtime which makes
the migration non-deterministic; change the sa.Column("embedding", Vector(...))
declaration to use a hardcoded integer literal for the vector dimension (replace
settings.EMBEDDING.VECTOR_DIMENSIONS) so the migration schema is immutable, and
update the migration docstring/comment near the top of this migration (e.g., the
module docstring for a1b2c3d4e5f6_initial_schema.py) to record the chosen
dimension and rationale; if you need to change dimensions later, create a
separate migration (e.g., a resize-embeddings migration) rather than reading
config at runtime.

sa.Column(
"created_at",
sa.DateTime(timezone=True),
Expand Down
2 changes: 1 addition & 1 deletion src/config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@

# OpenAI-compatible models that reject the `dimensions=` request parameter.
_EMBEDDING_KNOWN_REJECTING_MODELS: frozenset[str] = frozenset(
{"text-embedding-ada-002"}
{"text-embedding-ada-002", "mistral-embed"}
)
Comment on lines 30 to 32

@coderabbitai coderabbitai Bot Jun 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

🧩 Analysis chain

🌐 Web query:

Does Mistral's "mistral-embed" embedding model reject or ignore the dimensions parameter in API requests?

💡 Result:

Mistral’s embeddings API supports an output_dimension request field, not a dimensions field. The API documentation states that output_dimension is “the dimension of the output embeddings when feature available” and that “if not provided, a default output dimension will be used” [1]. The same field is present in the official Python SDK docs as output_dimension with the same description [2]. Therefore, if you send a dimensions parameter in the request for mistral-embed, there is no evidence in the official API spec that it is recognized; the documented parameter to control embedding size is output_dimension [1]. Also, third-party SDK docs note that output_dimension is only supported by certain models (giving “e.g.” examples), implying that for models where the feature isn’t available, requesting a different dimension won’t take effect as expected [3][4]. Practical guidance: use output_dimension (and only for models where Mistral documents it as supported) rather than dimensions when calling the embeddings endpoint for mistral-embed [1][2][3].

Citations:

Adjust Mistral embedding “dimensions” handling for mistral-embed

  • Mistral’s embeddings API documents output_dimension (and not a dimensions request field) to control embedding size, so excluding mistral-embed from sending dimensions is consistent with the official API spec to avoid unsupported-parameter behavior.
  • If the app supports user-specified embedding dimensions, map those to Mistral’s output_dimension (instead of only omitting dimensions) and add/extend coverage to assert the correct request field for mistral-embed.
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/config.py` around lines 30 - 32, The constant
_EMBEDDING_KNOWN_REJECTING_MODELS currently lists "mistral-embed" as a model
that rejects a "dimensions" field, but we must instead handle Mistral by mapping
any user-specified embedding dimensions to the Mistral request field
output_dimension (and not send dimensions). Update the embedding request
construction logic (where embeddings are built for models including
mistral-embed) to: detect model == "mistral-embed", translate any provided
dimensions value to output_dimension in the request payload, ensure the
dimensions field is omitted for that model, and add/extend unit/assertion
coverage to verify that requests for mistral-embed contain output_dimension (and
do not contain dimensions); reference the constant
_EMBEDDING_KNOWN_REJECTING_MODELS when locating related handling.



Expand Down