Fix the wrong search result of the pgvector (#514)

Signed-off-by: SimFG <bang.fu@zilliz.com>
zilliztech · Aug 12, 2023 · e242160 · RayceRossum · Aug 14, 2023 · e242160
1 parent 8f7d36a
commit e242160
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 5 deletions.
diff --git a/gptcache/__init__.py b/gptcache/__init__.py
@@ -1,5 +1,5 @@
 """gptcache version"""
-__version__ = "0.1.38"
+__version__ = "0.1.39"
 
 from gptcache.config import Config
 from gptcache.core import Cache

diff --git a/gptcache/manager/vector_data/pgvector.py b/gptcache/manager/vector_data/pgvector.py
@@ -149,6 +149,7 @@ def search(self, data: np.ndarray, top_k: int = -1):
             ).order_by(
                 similarity
             ).limit(top_k).all()
+            search_result = [(r[0].id, r[1]) for r in search_result]
 
         return search_result
 

diff --git a/tests/unit_tests/manager/test_pgvector.py b/tests/unit_tests/manager/test_pgvector.py
@@ -10,7 +10,7 @@
 class TestPgvector(unittest.TestCase):
     def test_normal(self):
         size = 1000
-        dim = 512
+        dim = 10
         top_k = 10
 
         url = os.getenv("POSTGRES_URL", "postgresql://postgres:postgres@localhost:5432/postgres")
@@ -31,10 +31,12 @@ def test_normal(self):
         self.assertEqual(len(db.search(data[0])), top_k)
         db.mul_add([VectorData(id=size, data=data[0])])
         ret = db.search(data[0])
-        self.assertIn(ret[0][1], [0, size])
-        self.assertIn(ret[1][1], [0, size])
+        print(ret)
+        self.assertIn(ret[0][0], [0, size])
+        self.assertIn(ret[1][0], [0, size])
         db.delete([0, 1, 2, 3, 4, 5, size])
         ret = db.search(data[0])
-        self.assertNotIn(ret[0][1], [0, size])
+        print(ret)
+        self.assertNotIn(ret[0][0], [0, size])
         db.rebuild()
         db.close()