From 897bbac699a69e068171018b5663b3b15c266f5a Mon Sep 17 00:00:00 2001 From: Boris Arzentar Date: Thu, 7 Nov 2024 11:36:31 +0100 Subject: [PATCH] fix: serialize UUID in pgvector data point payload --- .../databases/vector/pgvector/PGVectorAdapter.py | 4 ++-- .../{serialize_datetime.py => serialize_data.py} | 9 ++++++--- 2 files changed, 8 insertions(+), 5 deletions(-) rename cognee/infrastructure/databases/vector/pgvector/{serialize_datetime.py => serialize_data.py} (57%) diff --git a/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py b/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py index 321318043..235cc7745 100644 --- a/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py +++ b/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py @@ -8,7 +8,7 @@ from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker from cognee.infrastructure.engine import DataPoint -from .serialize_datetime import serialize_datetime +from .serialize_data import serialize_data from ..models.ScoredResult import ScoredResult from ..vector_db_interface import VectorDBInterface from ..embeddings.EmbeddingEngine import EmbeddingEngine @@ -113,7 +113,7 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface): PGVectorDataPoint( id=data_point.id, vector=data_vectors[data_index], - payload=serialize_datetime(data_point.model_dump()), + payload=serialize_data(data_point.model_dump()), ) for (data_index, data_point) in enumerate(data_points) ] diff --git a/cognee/infrastructure/databases/vector/pgvector/serialize_datetime.py b/cognee/infrastructure/databases/vector/pgvector/serialize_data.py similarity index 57% rename from cognee/infrastructure/databases/vector/pgvector/serialize_datetime.py rename to cognee/infrastructure/databases/vector/pgvector/serialize_data.py index 9cb979e2c..cdba1e928 100644 --- a/cognee/infrastructure/databases/vector/pgvector/serialize_datetime.py +++ b/cognee/infrastructure/databases/vector/pgvector/serialize_data.py @@ -1,12 +1,15 @@ from datetime import datetime +from uuid import UUID -def serialize_datetime(data): +def serialize_data(data): """Recursively convert datetime objects in dictionaries/lists to ISO format.""" if isinstance(data, dict): - return {key: serialize_datetime(value) for key, value in data.items()} + return {key: serialize_data(value) for key, value in data.items()} elif isinstance(data, list): - return [serialize_datetime(item) for item in data] + return [serialize_data(item) for item in data] elif isinstance(data, datetime): return data.isoformat() # Convert datetime to ISO 8601 string + elif isinstance(data, UUID): + return str(data) else: return data \ No newline at end of file