From 4a73505e23411baf10887420f1dd88ec1ac0406f Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Tue, 22 Oct 2024 10:29:35 +0200 Subject: [PATCH] refactor: Move create_db_and_tables module from vectors to pgvector Moved create_db_and_tables module to pgvector so it only handles pgvector related db creation Refactor #COG-170 --- cognee/api/v1/add/add.py | 6 +++--- cognee/api/v1/add/add_v2.py | 4 ++-- cognee/infrastructure/databases/vector/__init__.py | 1 - cognee/infrastructure/databases/vector/pgvector/__init__.py | 3 ++- .../databases/vector/{ => pgvector}/create_db_and_tables.py | 4 ++-- 5 files changed, 9 insertions(+), 9 deletions(-) rename cognee/infrastructure/databases/vector/{ => pgvector}/create_db_and_tables.py (81%) diff --git a/cognee/api/v1/add/add.py b/cognee/api/v1/add/add.py index b0b2fc8c6..10430ed8d 100644 --- a/cognee/api/v1/add/add.py +++ b/cognee/api/v1/add/add.py @@ -8,18 +8,18 @@ from cognee.infrastructure.files.storage import LocalStorage from cognee.modules.ingestion import get_matched_datasets, save_data_to_file from cognee.shared.utils import send_telemetry from cognee.base_config import get_base_config -from cognee.infrastructure.databases.relational import get_relational_engine, create_db_and_tables +from cognee.infrastructure.databases.relational import get_relational_engine from cognee.modules.users.methods import get_default_user from cognee.tasks.ingestion import get_dlt_destination from cognee.modules.users.permissions.methods import give_permission_on_document from cognee.modules.users.models import User from cognee.modules.data.methods import create_dataset from cognee.infrastructure.databases.relational import create_db_and_tables as create_relational_db_and_tables -from cognee.infrastructure.databases.vector import create_db_and_tables as create_vector_db_and_tables +from cognee.infrastructure.databases.vector.pgvector import create_db_and_tables as create_pgvector_db_and_tables async def add(data: Union[BinaryIO, List[BinaryIO], str, List[str]], dataset_name: str = "main_dataset", user: User = None): await create_relational_db_and_tables() - await create_vector_db_and_tables() + await create_pgvector_db_and_tables() if isinstance(data, str): if "data://" in data: diff --git a/cognee/api/v1/add/add_v2.py b/cognee/api/v1/add/add_v2.py index f32f470a9..4d43dd652 100644 --- a/cognee/api/v1/add/add_v2.py +++ b/cognee/api/v1/add/add_v2.py @@ -4,11 +4,11 @@ from cognee.modules.users.methods import get_default_user from cognee.modules.pipelines import run_tasks, Task from cognee.tasks.ingestion import save_data_to_storage, ingest_data from cognee.infrastructure.databases.relational import create_db_and_tables as create_relational_db_and_tables -from cognee.infrastructure.databases.vector import create_db_and_tables as create_vector_db_and_tables +from cognee.infrastructure.databases.vector.pgvector import create_db_and_tables as create_pgvector_db_and_tables async def add(data: Union[BinaryIO, list[BinaryIO], str, list[str]], dataset_name: str = "main_dataset", user: User = None): await create_relational_db_and_tables() - await create_vector_db_and_tables() + await create_pgvector_db_and_tables() if user is None: user = await get_default_user() diff --git a/cognee/infrastructure/databases/vector/__init__.py b/cognee/infrastructure/databases/vector/__init__.py index 02d13bb9c..604170f1d 100644 --- a/cognee/infrastructure/databases/vector/__init__.py +++ b/cognee/infrastructure/databases/vector/__init__.py @@ -4,4 +4,3 @@ from .models.CollectionConfig import CollectionConfig from .vector_db_interface import VectorDBInterface from .config import get_vectordb_config from .get_vector_engine import get_vector_engine -from .create_db_and_tables import create_db_and_tables diff --git a/cognee/infrastructure/databases/vector/pgvector/__init__.py b/cognee/infrastructure/databases/vector/pgvector/__init__.py index 84dc89113..130246a31 100644 --- a/cognee/infrastructure/databases/vector/pgvector/__init__.py +++ b/cognee/infrastructure/databases/vector/pgvector/__init__.py @@ -1 +1,2 @@ -from .PGVectorAdapter import PGVectorAdapter \ No newline at end of file +from .PGVectorAdapter import PGVectorAdapter +from .create_db_and_tables import create_db_and_tables \ No newline at end of file diff --git a/cognee/infrastructure/databases/vector/create_db_and_tables.py b/cognee/infrastructure/databases/vector/pgvector/create_db_and_tables.py similarity index 81% rename from cognee/infrastructure/databases/vector/create_db_and_tables.py rename to cognee/infrastructure/databases/vector/pgvector/create_db_and_tables.py index 21522db6b..1b09e873c 100644 --- a/cognee/infrastructure/databases/vector/create_db_and_tables.py +++ b/cognee/infrastructure/databases/vector/pgvector/create_db_and_tables.py @@ -1,5 +1,5 @@ -from ..relational.ModelBase import Base -from .get_vector_engine import get_vector_engine, get_vectordb_config +from ...relational.ModelBase import Base +from ..get_vector_engine import get_vector_engine, get_vectordb_config from sqlalchemy import text async def create_db_and_tables():