Merge pull request #1040 from airweave-ai/fix/embedders-vector-size

Fix: Embedder vector size

Author: marc-rutzou

backend/airweave/core/collection_service.py

@@ -44,6 +44,8 @@ async def _create(
         uow: UnitOfWork,
     ) -> schemas.Collection:
         """Create a new collection."""
+        from airweave.platform.destinations.collection_strategy import get_default_vector_size
+
         # Check if the collection already exists
         try:
             existing_collection = await crud.collection.get_by_readable_id(
@@ -57,20 +59,39 @@ async def _create(
                 status_code=400, detail="Collection with this readable_id already exists"
             )
 
-        collection = await crud.collection.create(db, obj_in=collection_in, ctx=ctx, uow=uow)
+        # Determine vector size and embedding model for this collection
+        vector_size = get_default_vector_size()
+
+        # Determine embedding model name based on vector size
+        from airweave.platform.destinations.collection_strategy import (
+            get_openai_embedding_model_for_vector_size,
+        )
+
+        try:
+            embedding_model_name = get_openai_embedding_model_for_vector_size(vector_size)
+        except ValueError:
+            # For non-OpenAI vector sizes (e.g., 384), use a generic name
+            embedding_model_name = "sentence-transformers/all-MiniLM-L6-v2"
+
+        # Add vector_size and embedding_model_name to collection data
+        collection_data = collection_in.model_dump()
+        collection_data["vector_size"] = vector_size
+        collection_data["embedding_model_name"] = embedding_model_name
+
+        collection = await crud.collection.create(db, obj_in=collection_data, ctx=ctx, uow=uow)
         await uow.session.flush()
 
-        # Create Qdrant destination with organization context
-        # Vector size is auto-detected based on embedding model configuration
+        # Create Qdrant destination with explicit vector size
         qdrant_destination = await QdrantDestination.create(
             credentials=None,  # Native Qdrant uses settings
             config=None,
             collection_id=collection.id,
             organization_id=ctx.organization.id,
+            vector_size=vector_size,
             logger=ctx.logger,
         )
 
-        # Setup the physical shared collection (auto-detects vector size)
+        # Setup the physical shared collection
         await qdrant_destination.setup_collection()
 
         return schemas.Collection.model_validate(collection, from_attributes=True)

backend/airweave/models/collection.py

@@ -2,7 +2,7 @@
 
 from typing import TYPE_CHECKING, List
 
-from sqlalchemy import String
+from sqlalchemy import Integer, String
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from airweave.models._base import OrganizationBase, UserMixin
@@ -19,6 +19,8 @@ class Collection(OrganizationBase, UserMixin):
 
     name: Mapped[str] = mapped_column(String, nullable=False)
     readable_id: Mapped[str] = mapped_column(String, nullable=False, unique=True)
+    vector_size: Mapped[int] = mapped_column(Integer, nullable=False)
+    embedding_model_name: Mapped[str] = mapped_column(String, nullable=False)
     # Status is now ephemeral - removed from database model
 
     # Relationships

backend/airweave/platform/destinations/collection_strategy.py

@@ -3,6 +3,7 @@
 All collections now use shared physical collections in Qdrant:
 - 384-dim vectors → airweave_shared_minilm_l6_v2 (local model)
 - 1536-dim vectors → airweave_shared_text_embedding_3_small (OpenAI)
+- 3072-dim vectors → airweave_shared_text_embedding_3_large (OpenAI)
 
 Tenant isolation is achieved via airweave_collection_id payload filtering.
 """
@@ -31,13 +32,38 @@ def get_physical_collection_name(vector_size: int | None = None) -> str:
     Returns:
         Physical collection name in Qdrant:
         - "airweave_shared_text_embedding_3_large" for 3072-dim vectors
+        - "airweave_shared_text_embedding_3_small" for 1536-dim vectors
         - "airweave_shared_minilm_l6_v2" for 384-dim vectors (default for other sizes)
     """
     if vector_size is None:
         vector_size = get_default_vector_size()
 
-    return (
-        "airweave_shared_text_embedding_3_large"
-        if vector_size == 3072
-        else "airweave_shared_minilm_l6_v2"
-    )
+    if vector_size == 3072:
+        return "airweave_shared_text_embedding_3_large"
+    elif vector_size == 1536:
+        return "airweave_shared_text_embedding_3_small"
+    else:
+        return "airweave_shared_minilm_l6_v2"
+
+
+def get_openai_embedding_model_for_vector_size(vector_size: int) -> str:
+    """Get OpenAI embedding model name for given vector dimensions.
+
+    Args:
+        vector_size: Vector dimensions (3072 or 1536)
+
+    Returns:
+        - "text-embedding-3-large" for 3072-dim
+        - "text-embedding-3-small" for 1536-dim
+
+    Raises:
+        ValueError: For vector sizes that don't use OpenAI models (e.g., 384 uses local model)
+    """
+    if vector_size == 3072:
+        return "text-embedding-3-large"
+    elif vector_size == 1536:
+        return "text-embedding-3-small"
+    else:
+        raise ValueError(
+            f"No OpenAI model for vector_size {vector_size}. Only 3072 and 1536 use OpenAI models."
+        )

backend/airweave/platform/embedders/openai.py

@@ -15,40 +15,64 @@
 
 
 class DenseEmbedder(BaseEmbedder):
-    """Singleton dense embedder using OpenAI text-embedding-3-large (3072 dims).
+    """OpenAI dense embedder with dynamic model selection (non-singleton).
+
+    IMPORTANT: No longer a singleton! Each collection may use different embedding models,
+    so we create fresh instances with the correct model for each sync/search operation.
 
     Features:
-    - Singleton shared across all syncs in pod
+    - Dynamic model selection based on vector_size (3072 or 1536)
     - Batch processing with OpenAI limits (2048 texts/request, 300K tokens/request)
     - 5 concurrent requests max
-    - Rate limiting with OpenAIRateLimiter singleton
+    - Rate limiting with OpenAIRateLimiter singleton (shared across instances)
     - Automatic retry on transient errors (via AsyncOpenAI client)
     - Fail-fast on any API errors (no silent failures)
     """
 
-    MODEL_NAME = "text-embedding-3-large"
-    VECTOR_DIMENSIONS = 3072
     MAX_TOKENS_PER_TEXT = 8192  # OpenAI limit per text
     MAX_BATCH_SIZE = 2048  # OpenAI limit per request
     MAX_TOKENS_PER_REQUEST = 300000  # OpenAI limit
     MAX_CONCURRENT_REQUESTS = 5
 
-    def __init__(self):
-        """Initialize OpenAI embedder (once per pod)."""
-        if self._initialized:
-            return
+    def __new__(cls, vector_size: int = None):
+        """Override singleton pattern from BaseEmbedder - create fresh instances."""
+        return object.__new__(cls)
+
+    def __init__(self, vector_size: int = None):
+        """Initialize OpenAI embedder for specific vector dimensions.
 
+        Args:
+            vector_size: Vector dimensions to determine model:
+                - 3072: text-embedding-3-large
+                - 1536: text-embedding-3-small
+                - None: defaults to 3072 (large model)
+        """
         if not settings.OPENAI_API_KEY:
             raise SyncFailureError("OPENAI_API_KEY required for dense embeddings")
 
+        # Fail-fast: vector_size should always be provided from collection
+        # Only allow None for backward compatibility, but warn
+        if vector_size is None:
+            # Fallback to large model but this shouldn't happen
+            self.MODEL_NAME = "text-embedding-3-large"
+            self.VECTOR_DIMENSIONS = 3072
+        else:
+            # Select model and dimensions based on vector_size
+            from airweave.platform.destinations.collection_strategy import (
+                get_openai_embedding_model_for_vector_size,
+            )
+
+            self.MODEL_NAME = get_openai_embedding_model_for_vector_size(vector_size)
+            self.VECTOR_DIMENSIONS = vector_size
+
+        # Create fresh client instance
         self._client = AsyncOpenAI(
             api_key=settings.OPENAI_API_KEY,
             timeout=1200.0,  # 20 min timeout for high concurrency
             max_retries=2,
         )
-        self._rate_limiter = OpenAIRateLimiter()  # Singleton
+        self._rate_limiter = OpenAIRateLimiter()  # This singleton is still OK (shared rate limit)
         self._tokenizer = tiktoken.get_encoding("cl100k_base")
-        self._initialized = True
 
     async def embed_many(self, texts: List[str], sync_context: SyncContext) -> List[List[float]]:
         """Embed batch of texts using OpenAI text-embedding-3-large.

backend/airweave/platform/sync/entity_pipeline.py

@@ -1167,9 +1167,10 @@ async def _embed_entities(
             sparse_texts.append(json.dumps(entity_dict, sort_keys=True))
 
         # Compute dense embeddings (always required)
+        # Create embedder with collection's vector_size (creates fresh instance)
         from airweave.platform.embedders import DenseEmbedder
 
-        dense_embedder = DenseEmbedder()
+        dense_embedder = DenseEmbedder(vector_size=sync_context.collection.vector_size)
         dense_embeddings = await dense_embedder.embed_many(dense_texts, sync_context)
 
         # Compute sparse embeddings (only if destination supports keyword index)

backend/airweave/platform/sync/factory.py

@@ -951,7 +951,7 @@ async def _get_integration_credential(
         return credential
 
     @classmethod
-    async def _create_destination_instances(
+    async def _create_destination_instances(  # noqa: C901
         cls,
         db: AsyncSession,
         sync: schemas.Sync,
@@ -997,12 +997,17 @@ async def _create_destination_instances(
                     destination_schema = schemas.Destination.model_validate(destination_model)
                     destination_class = resource_locator.get_destination(destination_schema)
 
+                    # Fail-fast: vector_size must be set
+                    if collection.vector_size is None:
+                        raise ValueError(f"Collection {collection.id} has no vector_size set.")
+
                     # Native Qdrant: no credentials (uses settings)
                     destination = await destination_class.create(
                         credentials=None,
                         config=None,
                         collection_id=collection.id,
                         organization_id=collection.organization_id,
+                        vector_size=collection.vector_size,
                         logger=logger,
                     )
 

backend/airweave/schemas/collection.py

@@ -166,6 +166,22 @@ class CollectionInDBBase(CollectionBase):
             "once the collection is created."
         ),
     )
+    vector_size: int = Field(
+        ...,
+        description=(
+            "Vector dimensions used by this collection. Determines which embedding model "
+            "is used: 3072 (text-embedding-3-large), 1536 (text-embedding-3-small), "
+            "or 384 (MiniLM-L6-v2)."
+        ),
+    )
+    embedding_model_name: str = Field(
+        ...,
+        description=(
+            "Name of the embedding model used for this collection "
+            "(e.g., 'text-embedding-3-large', 'text-embedding-3-small'). "
+            "This ensures queries use the same model as the indexed data."
+        ),
+    )
     created_at: datetime = Field(
         ...,
         description="Timestamp when the collection was created (ISO 8601 format).",

backend/airweave/search/defaults.yml

@@ -23,6 +23,16 @@ provider_models:
       name: "gpt-5-nano"
       tokenizer: "cl100k_base"
       context_window: 400000
+    embedding_large:
+      name: "text-embedding-3-large"
+      tokenizer: "cl100k_base"
+      dimensions: 3072
+      max_tokens: 8191
+    embedding_small:
+      name: "text-embedding-3-small"
+      tokenizer: "cl100k_base"
+      dimensions: 1536
+      max_tokens: 8191
     embedding:
       name: "text-embedding-3-large"
       tokenizer: "cl100k_base"