From 1a4fa0c3c074c9a0a0f1b7898f00401c5f91224a Mon Sep 17 00:00:00 2001
From: Javier Torres <javier@javiertorres.eu>
Date: Thu, 1 Aug 2024 10:12:53 +0200
Subject: [PATCH] Fix Matryohka check

---
 .../src/nucliadb/common/datamanagers/kb.py    | 28 +++++++++++--------
 1 file changed, 16 insertions(+), 12 deletions(-)

diff --git a/nucliadb/src/nucliadb/common/datamanagers/kb.py b/nucliadb/src/nucliadb/common/datamanagers/kb.py
index 5fb985dd7f..0a3bd96b5c 100644
--- 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)