Fix Matryohka check

nuclia · Aug 1, 2024 · 1a4fa0c · 1a4fa0c
1 parent 0f84ed8
commit 1a4fa0c
Showing 1 changed file with 16 additions and 12 deletions.
diff --git a/nucliadb/src/nucliadb/common/datamanagers/kb.py b/nucliadb/src/nucliadb/common/datamanagers/kb.py
@@ -111,18 +111,22 @@ async def get_matryoshka_vector_dimension(
     from . import vectorsets
 
     async for _, vs in vectorsets.iter(txn, kbid=kbid):
-        if vs.vectorset_index_config.vector_dimension in vs.matryoshka_dimensions:
-            return vs.vectorset_index_config.vector_dimension
-        else:
-            logger.error(
-                "KB has an invalid matryoshka dimension!",
-                extra={
-                    "kbid": kbid,
-                    "vector_dimension": vs.vectorset_index_config.vector_dimension,
-                    "matryoshka_dimensions": vs.matryoshka_dimensions,
-                },
-            )
-            return None
+        if (
+            len(vs.matryoshka_dimensions.matryoshka_dimensions) > 0
+            and vs.vectorset_index_config.vector_dimension
+        ):
+            if vs.vectorset_index_config.vector_dimension in vs.matryoshka_dimensions:
+                return vs.vectorset_index_config.vector_dimension
+            else:
+                logger.error(
+                    "KB has an invalid matryoshka dimension!",
+                    extra={
+                        "kbid": kbid,
+                        "vector_dimension": vs.vectorset_index_config.vector_dimension,
+                        "matryoshka_dimensions": vs.matryoshka_dimensions,
+                    },
+                )
+        return None
     else:
         # fallback for KBs that don't have vectorset
         model_metadata = await get_model_metadata(txn, kbid=kbid)