path: root/R2R/r2r/vecs/client.py

"""
Defines the 'Client' class

Importing from the `vecs.client` directly is not supported.
All public classes, enums, and functions are re-exported by the top level `vecs` module.
"""

from __future__ import annotations

import logging
import time
from typing import TYPE_CHECKING, List, Optional

import sqlalchemy
from deprecated import deprecated
from sqlalchemy import MetaData, create_engine, text
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import QueuePool

from .adapter import Adapter
from .exc import CollectionNotFound

if TYPE_CHECKING:
    from r2r.vecs.collection import Collection

logger = logging.getLogger(__name__)


class Client:
    """
    The `vecs.Client` class serves as an interface to a PostgreSQL database with pgvector support. It facilitates
    the creation, retrieval, listing and deletion of vector collections, while managing connections to the
    database.

    A `Client` instance represents a connection to a PostgreSQL database. This connection can be used to create
    and manipulate vector collections, where each collection is a group of vector records in a PostgreSQL table.

    The `vecs.Client` class can be also supports usage as a context manager to ensure the connection to the database
    is properly closed after operations, or used directly.

    Example usage:

        DB_CONNECTION = "postgresql://<user>:<password>@<host>:<port>/<db_name>"

        with vecs.create_client(DB_CONNECTION) as vx:
            # do some work
            pass

        # OR

        vx = vecs.create_client(DB_CONNECTION)
        # do some work
        vx.disconnect()
    """

    def __init__(
        self,
        connection_string: str,
        pool_size: int = 1,
        max_retries: int = 3,
        retry_delay: int = 1,
    ):
        self.engine = create_engine(
            connection_string,
            pool_size=pool_size,
            poolclass=QueuePool,
            pool_recycle=300,  # Recycle connections after 5 min
        )
        self.meta = MetaData(schema="vecs")
        self.Session = sessionmaker(self.engine)
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        self.vector_version: Optional[str] = None
        self._initialize_database()

    def _initialize_database(self):
        retries = 0
        error = None
        while retries < self.max_retries:
            try:
                with self.Session() as sess:
                    with sess.begin():
                        self._create_schema(sess)
                        self._create_extension(sess)
                        self._get_vector_version(sess)
                return
            except Exception as e:
                logger.warning(
                    f"Database connection error: {str(e)}. Retrying in {self.retry_delay} seconds..."
                )
                retries += 1
                time.sleep(self.retry_delay)
                error = e

        error_message = f"Failed to initialize database after {self.max_retries} retries with error: {str(error)}"
        logger.error(error_message)
        raise RuntimeError(error_message)

    def _create_schema(self, sess):
        try:
            sess.execute(text("CREATE SCHEMA IF NOT EXISTS vecs;"))
        except Exception as e:
            logger.warning(f"Failed to create schema: {str(e)}")

    def _create_extension(self, sess):
        try:
            sess.execute(text("CREATE EXTENSION IF NOT EXISTS vector;"))
            sess.execute(text("CREATE EXTENSION IF NOT EXISTS pg_trgm;"))
            sess.execute(text("CREATE EXTENSION IF NOT EXISTS fuzzystrmatch;"))
        except Exception as e:
            logger.warning(f"Failed to create extension: {str(e)}")

    def _get_vector_version(self, sess):
        try:
            self.vector_version = sess.execute(
                text(
                    "SELECT installed_version FROM pg_available_extensions WHERE name = 'vector' LIMIT 1;"
                )
            ).scalar_one()
        except sqlalchemy.exc.InternalError as e:
            logger.error(f"Failed with internal alchemy error: {str(e)}")

            import psycopg2

            if isinstance(e.orig, psycopg2.errors.InFailedSqlTransaction):
                sess.rollback()
                self.vector_version = sess.execute(
                    text(
                        "SELECT installed_version FROM pg_available_extensions WHERE name = 'vector' LIMIT 1;"
                    )
                ).scalar_one()
            else:
                raise e
        except Exception as e:
            logger.error(f"Failed to retrieve vector version: {str(e)}")
            raise e

    def _supports_hnsw(self):
        return (
            not self.vector_version.startswith("0.4")
            and not self.vector_version.startswith("0.3")
            and not self.vector_version.startswith("0.2")
            and not self.vector_version.startswith("0.1")
            and not self.vector_version.startswith("0.0")
        )

    def get_or_create_collection(
        self,
        name: str,
        *,
        dimension: Optional[int] = None,
        adapter: Optional[Adapter] = None,
    ) -> Collection:
        """
        Get a vector collection by name, or create it if no collection with
        *name* exists.

        Args:
            name (str): The name of the collection.

        Keyword Args:
            dimension (int): The dimensionality of the vectors in the collection.
            pipeline (int): The dimensionality of the vectors in the collection.

        Returns:
            Collection: The created collection.

        Raises:
            CollectionAlreadyExists: If a collection with the same name already exists
        """
        from r2r.vecs.collection import Collection

        adapter_dimension = adapter.exported_dimension if adapter else None

        collection = Collection(
            name=name,
            dimension=dimension or adapter_dimension,  # type: ignore
            client=self,
            adapter=adapter,
        )

        return collection._create_if_not_exists()

    @deprecated("use Client.get_or_create_collection")
    def create_collection(self, name: str, dimension: int) -> Collection:
        """
        Create a new vector collection.

        Args:
            name (str): The name of the collection.
            dimension (int): The dimensionality of the vectors in the collection.

        Returns:
            Collection: The created collection.

        Raises:
            CollectionAlreadyExists: If a collection with the same name already exists
        """
        from r2r.vecs.collection import Collection

        return Collection(name, dimension, self)._create()

    @deprecated("use Client.get_or_create_collection")
    def get_collection(self, name: str) -> Collection:
        """
        Retrieve an existing vector collection.

        Args:
            name (str): The name of the collection.

        Returns:
            Collection: The retrieved collection.

        Raises:
            CollectionNotFound: If no collection with the given name exists.
        """
        from r2r.vecs.collection import Collection

        query = text(
            f"""
        select
            relname as table_name,
            atttypmod as embedding_dim
        from
            pg_class pc
            join pg_attribute pa
                on pc.oid = pa.attrelid
        where
            pc.relnamespace = 'vecs'::regnamespace
            and pc.relkind = 'r'
            and pa.attname = 'vec'
            and not pc.relname ^@ '_'
            and pc.relname = :name
        """
        ).bindparams(name=name)
        with self.Session() as sess:
            query_result = sess.execute(query).fetchone()

            if query_result is None:
                raise CollectionNotFound(
                    "No collection found with requested name"
                )

            name, dimension = query_result
            return Collection(
                name,
                dimension,
                self,
            )

    def list_collections(self) -> List["Collection"]:
        """
        List all vector collections.

        Returns:
            list[Collection]: A list of all collections.
        """
        from r2r.vecs.collection import Collection

        return Collection._list_collections(self)

    def delete_collection(self, name: str) -> None:
        """
        Delete a vector collection.

        If no collection with requested name exists, does nothing.

        Args:
            name (str): The name of the collection.

        Returns:
            None
        """
        from r2r.vecs.collection import Collection

        Collection(name, -1, self)._drop()
        return

    def disconnect(self) -> None:
        """
        Disconnect the client from the database.

        Returns:
            None
        """
        self.engine.dispose()
        logger.info("Disconnected from the database.")
        return

    def __enter__(self) -> "Client":
        """
        Enable use of the 'with' statement.

        Returns:
            Client: The current instance of the Client.
        """

        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        Disconnect the client on exiting the 'with' statement context.

        Args:
            exc_type: The exception type, if any.
            exc_val: The exception value, if any.
            exc_tb: The traceback, if any.

        Returns:
            None
        """
        self.disconnect()
        return