diff --git a/mem0/vector_stores/pgvector.py b/mem0/vector_stores/pgvector.py index d36cab1cc..454313886 100644 --- a/mem0/vector_stores/pgvector.py +++ b/mem0/vector_stores/pgvector.py @@ -459,13 +459,13 @@ class PGVector(VectorStoreBase): self._ensure_collection() with self._get_cursor() as cur: cur.execute( - sql.SQL("SELECT id, vector, payload FROM {} WHERE id = %s").format(self._col()), + sql.SQL("SELECT id, payload FROM {} WHERE id = %s").format(self._col()), (vector_id,), ) result = cur.fetchone() if not result: return None - return OutputData(id=str(result[0]), score=None, payload=result[2]) + return OutputData(id=str(result[0]), score=None, payload=result[1]) def list_cols(self) -> List[str]: """ @@ -528,7 +528,7 @@ class PGVector(VectorStoreBase): with self._get_cursor() as cur: cur.execute( sql.SQL(""" - SELECT id, vector, payload + SELECT id, payload FROM {} {} LIMIT %s @@ -536,7 +536,7 @@ class PGVector(VectorStoreBase): (*filter_params, top_k), ) results = cur.fetchall() - return [[OutputData(id=str(r[0]), score=None, payload=r[2]) for r in results]] + return [[OutputData(id=str(r[0]), score=None, payload=r[1]) for r in results]] def __del__(self) -> None: """ diff --git a/tests/vector_stores/test_pgvector.py b/tests/vector_stores/test_pgvector.py index d3ac284bf..f4a3351af 100644 --- a/tests/vector_stores/test_pgvector.py +++ b/tests/vector_stores/test_pgvector.py @@ -711,7 +711,7 @@ class TestPGVector(unittest.TestCase): mock_get_cursor.return_value.__exit__.return_value = None self.mock_cursor.fetchall.return_value = [] # No existing collections - self.mock_cursor.fetchone.return_value = (self.test_ids[0], [0.1, 0.2, 0.3], {"key": "value1"}) + self.mock_cursor.fetchone.return_value = (self.test_ids[0], {"key": "value1"}) pgvector = PGVector( dbname="test_db", @@ -734,7 +734,7 @@ class TestPGVector(unittest.TestCase): # Verify get query was executed get_calls = [call for call in self.mock_cursor.execute.call_args_list - if "SELECT id, vector, payload" in str(call)] + if "SELECT id, payload" in str(call)] self.assertTrue(len(get_calls) > 0) # Verify result @@ -756,7 +756,7 @@ class TestPGVector(unittest.TestCase): mock_get_cursor.return_value.__exit__.return_value = None self.mock_cursor.fetchall.return_value = [] # No existing collections - self.mock_cursor.fetchone.return_value = (self.test_ids[0], [0.1, 0.2, 0.3], {"key": "value1"}) + self.mock_cursor.fetchone.return_value = (self.test_ids[0], {"key": "value1"}) pgvector = PGVector( dbname="test_db", @@ -779,7 +779,7 @@ class TestPGVector(unittest.TestCase): # Verify get query was executed get_calls = [call for call in self.mock_cursor.execute.call_args_list - if "SELECT id, vector, payload" in str(call)] + if "SELECT id, payload" in str(call)] self.assertTrue(len(get_calls) > 0) # Verify result @@ -1050,8 +1050,8 @@ class TestPGVector(unittest.TestCase): mock_get_cursor.return_value.__exit__.return_value = None self.mock_cursor.fetchall.return_value = [ - (self.test_ids[0], [0.1, 0.2, 0.3], {"key": "value1"}), - (self.test_ids[1], [0.4, 0.5, 0.6], {"key": "value2"}), + (self.test_ids[0], {"key": "value1"}), + (self.test_ids[1], {"key": "value2"}), ] pgvector = PGVector( @@ -1075,7 +1075,7 @@ class TestPGVector(unittest.TestCase): # Verify list query was executed list_calls = [call for call in self.mock_cursor.execute.call_args_list - if "SELECT id, vector, payload" in str(call)] + if "SELECT id, payload" in str(call)] self.assertTrue(len(list_calls) > 0) # Verify result @@ -1098,8 +1098,8 @@ class TestPGVector(unittest.TestCase): mock_get_cursor.return_value.__exit__.return_value = None self.mock_cursor.fetchall.return_value = [ - (self.test_ids[0], [0.1, 0.2, 0.3], {"key": "value1"}), - (self.test_ids[1], [0.4, 0.5, 0.6], {"key": "value2"}), + (self.test_ids[0], {"key": "value1"}), + (self.test_ids[1], {"key": "value2"}), ] pgvector = PGVector( @@ -1123,7 +1123,7 @@ class TestPGVector(unittest.TestCase): # Verify list query was executed list_calls = [call for call in self.mock_cursor.execute.call_args_list - if "SELECT id, vector, payload" in str(call)] + if "SELECT id, payload" in str(call)] self.assertTrue(len(list_calls) > 0) # Verify result @@ -1440,7 +1440,7 @@ class TestPGVector(unittest.TestCase): mock_get_cursor.return_value.__exit__.return_value = None self.mock_cursor.fetchall.return_value = [ - (self.test_ids[0], [0.1, 0.2, 0.3], {"user_id": "alice", "agent_id": "agent1"}), + (self.test_ids[0], {"user_id": "alice", "agent_id": "agent1"}), ] pgvector = PGVector( @@ -1465,7 +1465,7 @@ class TestPGVector(unittest.TestCase): # Verify list query was executed with filters list_calls = [call for call in self.mock_cursor.execute.call_args_list - if "SELECT id, vector, payload" in str(call) and "WHERE" in str(call)] + if "SELECT id, payload" in str(call) and "WHERE" in str(call)] self.assertTrue(len(list_calls) > 0) # Verify results @@ -1489,7 +1489,7 @@ class TestPGVector(unittest.TestCase): mock_get_cursor.return_value.__exit__.return_value = None self.mock_cursor.fetchall.return_value = [ - (self.test_ids[0], [0.1, 0.2, 0.3], {"user_id": "alice", "agent_id": "agent1"}), + (self.test_ids[0], {"user_id": "alice", "agent_id": "agent1"}), ] pgvector = PGVector( @@ -1514,7 +1514,7 @@ class TestPGVector(unittest.TestCase): # Verify list query was executed with filters list_calls = [call for call in self.mock_cursor.execute.call_args_list - if "SELECT id, vector, payload" in str(call) and "WHERE" in str(call)] + if "SELECT id, payload" in str(call) and "WHERE" in str(call)] self.assertTrue(len(list_calls) > 0) # Verify results @@ -1538,7 +1538,7 @@ class TestPGVector(unittest.TestCase): mock_get_cursor.return_value.__exit__.return_value = None self.mock_cursor.fetchall.return_value = [ - (self.test_ids[0], [0.1, 0.2, 0.3], {"user_id": "alice"}), + (self.test_ids[0], {"user_id": "alice"}), ] pgvector = PGVector( @@ -1563,7 +1563,7 @@ class TestPGVector(unittest.TestCase): # Verify list query was executed with single filter list_calls = [call for call in self.mock_cursor.execute.call_args_list - if "SELECT id, vector, payload" in str(call) and "WHERE" in str(call)] + if "SELECT id, payload" in str(call) and "WHERE" in str(call)] self.assertTrue(len(list_calls) > 0) # Verify results @@ -1586,7 +1586,7 @@ class TestPGVector(unittest.TestCase): mock_get_cursor.return_value.__exit__.return_value = None self.mock_cursor.fetchall.return_value = [ - (self.test_ids[0], [0.1, 0.2, 0.3], {"user_id": "alice"}), + (self.test_ids[0], {"user_id": "alice"}), ] pgvector = PGVector( @@ -1611,7 +1611,7 @@ class TestPGVector(unittest.TestCase): # Verify list query was executed with single filter list_calls = [call for call in self.mock_cursor.execute.call_args_list - if "SELECT id, vector, payload" in str(call) and "WHERE" in str(call)] + if "SELECT id, payload" in str(call) and "WHERE" in str(call)] self.assertTrue(len(list_calls) > 0) # Verify results @@ -1634,8 +1634,8 @@ class TestPGVector(unittest.TestCase): mock_get_cursor.return_value.__exit__.return_value = None self.mock_cursor.fetchall.return_value = [ - (self.test_ids[0], [0.1, 0.2, 0.3], {"key": "value1"}), - (self.test_ids[1], [0.4, 0.5, 0.6], {"key": "value2"}), + (self.test_ids[0], {"key": "value1"}), + (self.test_ids[1], {"key": "value2"}), ] pgvector = PGVector( @@ -1659,7 +1659,7 @@ class TestPGVector(unittest.TestCase): # Verify list query was executed without WHERE clause list_calls = [call for call in self.mock_cursor.execute.call_args_list - if "SELECT id, vector, payload" in str(call) and "WHERE" not in str(call)] + if "SELECT id, payload" in str(call) and "WHERE" not in str(call)] self.assertTrue(len(list_calls) > 0) # Verify results @@ -1682,8 +1682,8 @@ class TestPGVector(unittest.TestCase): mock_get_cursor.return_value.__exit__.return_value = None self.mock_cursor.fetchall.return_value = [ - (self.test_ids[0], [0.1, 0.2, 0.3], {"key": "value1"}), - (self.test_ids[1], [0.4, 0.5, 0.6], {"key": "value2"}), + (self.test_ids[0], {"key": "value1"}), + (self.test_ids[1], {"key": "value2"}), ] pgvector = PGVector( @@ -1707,7 +1707,7 @@ class TestPGVector(unittest.TestCase): # Verify list query was executed without WHERE clause list_calls = [call for call in self.mock_cursor.execute.call_args_list - if "SELECT id, vector, payload" in str(call) and "WHERE" not in str(call)] + if "SELECT id, payload" in str(call) and "WHERE" not in str(call)] self.assertTrue(len(list_calls) > 0) # Verify results