Repository: topoteretes/cognee
File: cognee/infrastructure/databases/vector/pgvector/PGVectorDatasetDatabaseHandler.py (Python)

Code before:
from uuid import UUID
from typing import Optional

from cognee.infrastructure.databases.vector.create_vector_engine import create_vector_engine
from cognee.modules.users.models import User
from cognee.modules.users.models import DatasetDatabase
from cognee.infrastructure.databases.vector import get_vectordb_config
from cognee.infrastructure.databases.dataset_database_handler import DatasetDatabaseHandlerInterface


class PGVectorDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
    """
    Handler for interacting with PGVector Dataset databases.
    """

    @classmethod
    async def _create_pg_database(cls, vector_config):
        """
        Create the necessary Postgres database, and the PGVector extension on it.
        This is defined here because the creation needs the latest vector config,
        which is not yet saved in the vector config context variable here.
        """
        from cognee.infrastructure.databases.relational.create_relational_engine import (
            create_relational_engine,
        )

        from cognee.infrastructure.databases.vector.create_vector_engine import create_vector_engine
        from sqlalchemy import text

        pg_relational_engine = create_relational_engine(
            db_path="",
            db_host=vector_config["vector_db_url"],
            db_name=vector_config["vector_db_name"],
            db_port=vector_config["vector_db_port"],
            db_username=vector_config["vector_db_username"],
            db_password=vector_config["vector_db_password"],
            db_provider="postgres",
        )
        await pg_relational_engine.create_database()

        vector_engine = create_vector_engine(**vector_config)
        async with vector_engine.engine.begin() as connection:

Review comment (refactor):
Can we move this as a utility function for the PGVector adapter? So we can use it in other dataset database handlers in the future

Code after:
from uuid import UUID
from typing import Optional

from cognee.infrastructure.databases.vector.create_vector_engine import create_vector_engine
from cognee.infrastructure.databases.vector.pgvector.create_db_and_tables import delete_pg_database
from cognee.modules.users.models import User
from cognee.modules.users.models import DatasetDatabase
from cognee.infrastructure.databases.vector import get_vectordb_config
from cognee.infrastructure.databases.dataset_database_handler import DatasetDatabaseHandlerInterface


class PGVectorDatasetDatabaseHandler(DatasetDatabaseHandlerInterface):
    """
    Handler for interacting with PGVector Dataset databases.
    """

    @classmethod
    async def create_dataset(cls, dataset_id: Optional[UUID], user: Optional[User]) -> dict:
        vector_config = get_vectordb_config()

        if vector_config.vector_db_provider != "pgvector":
            raise ValueError(
                "PGVectorDatasetDatabaseHandler can only be used with PGVector vector database provider."
            )

        vector_db_name = f"{dataset_id}"

        new_vector_config = {
            "vector_database_provider": vector_config.vector_db_provider,
            "vector_database_url": vector_config.vector_db_url,
            "vector_database_name": vector_db_name,
            "vector_database_connection_info": {
                "port": vector_config.vector_db_port,
                "host": vector_config.vector_db_host,
            },
            "vector_dataset_database_handler": "pgvector",
        }

        from .create_db_and_tables import create_pg_database

        await create_pg_database(
            {