diff --git a/scripts/tests/test-retrieval-server.py b/scripts/tests/test-retrieval-server.py index 9028a83..8faf6c0 100644 --- a/scripts/tests/test-retrieval-server.py +++ b/scripts/tests/test-retrieval-server.py @@ -403,6 +403,25 @@ class RetrievalServerTests(unittest.TestCase): self.assertEqual(MODULE.reciprocal_rank_fusion([1, 2, 3], [4, 2, 3]), [2, 3, 1, 4]) self.assertEqual(MODULE.reciprocal_rank_fusion([], [4, 2]), [4, 2]) + def test_dense_candidates_converts_vector_keys_for_sqlite(self): + class VectorKey: + def __int__(self): + return 9 + + with TemporaryDirectory() as temporary: + root = Path(temporary) + index = MODULE.Index(root, root / 'data') + index.vector_index = SimpleNamespace(search=lambda vector, count: [SimpleNamespace(key=VectorKey())]) + usearch, usearch_index = fake_usearch() + try: + with patch.dict('sys.modules', numpy=fake_numpy(), usearch=usearch, + **{'usearch.index': usearch_index}): + candidates = index.dense_candidates(embedding()) + self.assertEqual(candidates, [9]) + self.assertIs(type(candidates[0]), int) + finally: + index.db.close() + def test_hybrid_candidates_and_document_metadata(self): with TemporaryDirectory() as temporary: root = Path(temporary) diff --git a/shared/retrieval/server.py b/shared/retrieval/server.py index 525aefb..9a8277d 100644 --- a/shared/retrieval/server.py +++ b/shared/retrieval/server.py @@ -497,7 +497,7 @@ class Index: vectors = np.frombuffer(b"".join(row[1] for row in rows), dtype=np.float32).reshape(-1, EMBEDDING_DIMENSIONS) index.add([row[0] for row in rows], vectors) self.vector_index = index - return [match.key for match in self.vector_index.search(vector, DENSE_CANDIDATES)] + return [int(match.key) for match in self.vector_index.search(vector, DENSE_CANDIDATES)] def cache_key(self, query: str, top_k: int): identity = json.dumps([