langchain/libs/community/langchain_community/vectorstores/timescalevector.py

"""VectorStore wrapper around a Postgres-TimescaleVector database."""
from __future__ import annotations

import enum
import logging
import uuid
from datetime import timedelta
from typing import (
    TYPE_CHECKING,
    Any,
    Callable,
    Dict,
    Iterable,
    List,
    Optional,
    Tuple,
    Type,
    Union,
)

from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.utils import get_from_dict_or_env
from langchain_core.vectorstores import VectorStore

from langchain_community.vectorstores.utils import DistanceStrategy

if TYPE_CHECKING:
    from timescale_vector import Predicates


DEFAULT_DISTANCE_STRATEGY = DistanceStrategy.COSINE

ADA_TOKEN_COUNT = 1536

_LANGCHAIN_DEFAULT_COLLECTION_NAME = "langchain_store"


class TimescaleVector(VectorStore):
    """Timescale Postgres vector store

    To use, you should have the ``timescale_vector`` python package installed.

    Args:
        service_url: Service url on timescale cloud.
        embedding: Any embedding function implementing
            `langchain.embeddings.base.Embeddings` interface.
        collection_name: The name of the collection to use. (default: langchain_store)
            This will become the table name used for the collection.
        distance_strategy: The distance strategy to use. (default: COSINE)
        pre_delete_collection: If True, will delete the collection if it exists.
            (default: False). Useful for testing.

    Example:
        .. code-block:: python

            from langchain_community.vectorstores import TimescaleVector
            from langchain_community.embeddings.openai import OpenAIEmbeddings

            SERVICE_URL = "postgres://tsdbadmin:<password>@<id>.tsdb.cloud.timescale.com:<port>/tsdb?sslmode=require"
            COLLECTION_NAME = "state_of_the_union_test"
            embeddings = OpenAIEmbeddings()
            vectorestore = TimescaleVector.from_documents(
                embedding=embeddings,
                documents=docs,
                collection_name=COLLECTION_NAME,
                service_url=SERVICE_URL,
            )
    """  # noqa: E501

    def __init__(
        self,
        service_url: str,
        embedding: Embeddings,
        collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME,
        num_dimensions: int = ADA_TOKEN_COUNT,
        distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY,
        pre_delete_collection: bool = False,
        logger: Optional[logging.Logger] = None,
        relevance_score_fn: Optional[Callable[[float], float]] = None,
        time_partition_interval: Optional[timedelta] = None,
        **kwargs: Any,
    ) -> None:
        try:
            from timescale_vector import client
        except ImportError:
            raise ImportError(
                "Could not import timescale_vector python package. "
                "Please install it with `pip install timescale-vector`."
            )

        self.service_url = service_url
        self.embedding = embedding
        self.collection_name = collection_name
        self.num_dimensions = num_dimensions
        self._distance_strategy = distance_strategy
        self.pre_delete_collection = pre_delete_collection
        self.logger = logger or logging.getLogger(__name__)
        self.override_relevance_score_fn = relevance_score_fn
        self._time_partition_interval = time_partition_interval
        self.sync_client = client.Sync(
            self.service_url,
            self.collection_name,
            self.num_dimensions,
            self._distance_strategy.value.lower(),
            time_partition_interval=self._time_partition_interval,
            **kwargs,
        )
        self.async_client = client.Async(
            self.service_url,
            self.collection_name,
            self.num_dimensions,
            self._distance_strategy.value.lower(),
            time_partition_interval=self._time_partition_interval,
            **kwargs,
        )
        self.__post_init__()

    def __post_init__(
        self,
    ) -> None:
        """
        Initialize the store.
        """
        self.sync_client.create_tables()
        if self.pre_delete_collection:
            self.sync_client.delete_all()

    @property
    def embeddings(self) -> Embeddings:
        return self.embedding

    def drop_tables(self) -> None:
        self.sync_client.drop_table()

    @classmethod
    def __from(
        cls,
        texts: List[str],
        embeddings: List[List[float]],
        embedding: Embeddings,
        metadatas: Optional[List[dict]] = None,
        ids: Optional[List[str]] = None,
        collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME,
        distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY,
        service_url: Optional[str] = None,
        pre_delete_collection: bool = False,
        **kwargs: Any,
    ) -> TimescaleVector:
        num_dimensions = len(embeddings[0])

        if ids is None:
            ids = [str(uuid.uuid4()) for _ in texts]

        if not metadatas:
            metadatas = [{} for _ in texts]

        if service_url is None:
            service_url = cls.get_service_url(kwargs)

        store = cls(
            service_url=service_url,
            num_dimensions=num_dimensions,
            collection_name=collection_name,
            embedding=embedding,
            distance_strategy=distance_strategy,
            pre_delete_collection=pre_delete_collection,
            **kwargs,
        )

        store.add_embeddings(
            texts=texts, embeddings=embeddings, metadatas=metadatas, ids=ids, **kwargs
        )

        return store

    @classmethod
    async def __afrom(
        cls,
        texts: List[str],
        embeddings: List[List[float]],
        embedding: Embeddings,
        metadatas: Optional[List[dict]] = None,
        ids: Optional[List[str]] = None,
        collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME,
        distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY,
        service_url: Optional[str] = None,
        pre_delete_collection: bool = False,
        **kwargs: Any,
    ) -> TimescaleVector:
        num_dimensions = len(embeddings[0])

        if ids is None:
            ids = [str(uuid.uuid4()) for _ in texts]

        if not metadatas:
            metadatas = [{} for _ in texts]

        if service_url is None:
            service_url = cls.get_service_url(kwargs)

        store = cls(
            service_url=service_url,
            num_dimensions=num_dimensions,
            collection_name=collection_name,
            embedding=embedding,
            distance_strategy=distance_strategy,
            pre_delete_collection=pre_delete_collection,
            **kwargs,
        )

        await store.aadd_embeddings(
            texts=texts, embeddings=embeddings, metadatas=metadatas, ids=ids, **kwargs
        )

        return store

    def add_embeddings(
        self,
        texts: Iterable[str],
        embeddings: List[List[float]],
        metadatas: Optional[List[dict]] = None,
        ids: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> List[str]:
        """Add embeddings to the vectorstore.

        Args:
            texts: Iterable of strings to add to the vectorstore.
            embeddings: List of list of embedding vectors.
            metadatas: List of metadatas associated with the texts.
            kwargs: vectorstore specific parameters
        """
        if ids is None:
            ids = [str(uuid.uuid4()) for _ in texts]

        if not metadatas:
            metadatas = [{} for _ in texts]

        records = list(zip(ids, metadatas, texts, embeddings))
        self.sync_client.upsert(records)

        return ids

    async def aadd_embeddings(
        self,
        texts: Iterable[str],
        embeddings: List[List[float]],
        metadatas: Optional[List[dict]] = None,
        ids: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> List[str]:
        """Add embeddings to the vectorstore.

        Args:
            texts: Iterable of strings to add to the vectorstore.
            embeddings: List of list of embedding vectors.
            metadatas: List of metadatas associated with the texts.
            kwargs: vectorstore specific parameters
        """
        if ids is None:
            ids = [str(uuid.uuid4()) for _ in texts]

        if not metadatas:
            metadatas = [{} for _ in texts]

        records = list(zip(ids, metadatas, texts, embeddings))
        await self.async_client.upsert(records)

        return ids

    def add_texts(
        self,
        texts: Iterable[str],
        metadatas: Optional[List[dict]] = None,
        ids: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> List[str]:
        """Run more texts through the embeddings and add to the vectorstore.

        Args:
            texts: Iterable of strings to add to the vectorstore.
            metadatas: Optional list of metadatas associated with the texts.
            kwargs: vectorstore specific parameters

        Returns:
            List of ids from adding the texts into the vectorstore.
        """
        embeddings = self.embedding.embed_documents(list(texts))
        return self.add_embeddings(
            texts=texts, embeddings=embeddings, metadatas=metadatas, ids=ids, **kwargs
        )

    async def aadd_texts(
        self,
        texts: Iterable[str],
        metadatas: Optional[List[dict]] = None,
        ids: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> List[str]:
        """Run more texts through the embeddings and add to the vectorstore.

        Args:
            texts: Iterable of strings to add to the vectorstore.
            metadatas: Optional list of metadatas associated with the texts.
            kwargs: vectorstore specific parameters

        Returns:
            List of ids from adding the texts into the vectorstore.
        """
        embeddings = self.embedding.embed_documents(list(texts))
        return await self.aadd_embeddings(
            texts=texts, embeddings=embeddings, metadatas=metadatas, ids=ids, **kwargs
        )

    def _embed_query(self, query: str) -> Optional[List[float]]:
        # an empty query should not be embedded
        if query is None or query == "" or query.isspace():
            return None
        else:
            return self.embedding.embed_query(query)

    def similarity_search(
        self,
        query: str,
        k: int = 4,
        filter: Optional[Union[dict, list]] = None,
        predicates: Optional[Predicates] = None,
        **kwargs: Any,
    ) -> List[Document]:
        """Run similarity search with TimescaleVector with distance.

        Args:
            query (str): Query text to search for.
            k (int): Number of results to return. Defaults to 4.
            filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.

        Returns:
            List of Documents most similar to the query.
        """
        embedding = self._embed_query(query)
        return self.similarity_search_by_vector(
            embedding=embedding,
            k=k,
            filter=filter,
            predicates=predicates,
            **kwargs,
        )

    async def asimilarity_search(
        self,
        query: str,
        k: int = 4,
        filter: Optional[Union[dict, list]] = None,
        predicates: Optional[Predicates] = None,
        **kwargs: Any,
    ) -> List[Document]:
        """Run similarity search with TimescaleVector with distance.

        Args:
            query (str): Query text to search for.
            k (int): Number of results to return. Defaults to 4.
            filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.

        Returns:
            List of Documents most similar to the query.
        """
        embedding = self._embed_query(query)
        return await self.asimilarity_search_by_vector(
            embedding=embedding,
            k=k,
            filter=filter,
            predicates=predicates,
            **kwargs,
        )

    def similarity_search_with_score(
        self,
        query: str,
        k: int = 4,
        filter: Optional[Union[dict, list]] = None,
        predicates: Optional[Predicates] = None,
        **kwargs: Any,
    ) -> List[Tuple[Document, float]]:
        """Return docs most similar to query.

        Args:
            query: Text to look up documents similar to.
            k: Number of Documents to return. Defaults to 4.
            filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.

        Returns:
            List of Documents most similar to the query and score for each
        """
        embedding = self._embed_query(query)
        docs = self.similarity_search_with_score_by_vector(
            embedding=embedding,
            k=k,
            filter=filter,
            predicates=predicates,
            **kwargs,
        )
        return docs

    async def asimilarity_search_with_score(
        self,
        query: str,
        k: int = 4,
        filter: Optional[Union[dict, list]] = None,
        predicates: Optional[Predicates] = None,
        **kwargs: Any,
    ) -> List[Tuple[Document, float]]:
        """Return docs most similar to query.

        Args:
            query: Text to look up documents similar to.
            k: Number of Documents to return. Defaults to 4.
            filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.

        Returns:
            List of Documents most similar to the query and score for each
        """

        embedding = self._embed_query(query)
        return await self.asimilarity_search_with_score_by_vector(
            embedding=embedding,
            k=k,
            filter=filter,
            predicates=predicates,
            **kwargs,
        )

    def date_to_range_filter(self, **kwargs: Any) -> Any:
        constructor_args = {
            key: kwargs[key]
            for key in [
                "start_date",
                "end_date",
                "time_delta",
                "start_inclusive",
                "end_inclusive",
            ]
            if key in kwargs
        }
        if not constructor_args or len(constructor_args) == 0:
            return None

        try:
            from timescale_vector import client
        except ImportError:
            raise ImportError(
                "Could not import timescale_vector python package. "
                "Please install it with `pip install timescale-vector`."
            )
        return client.UUIDTimeRange(**constructor_args)

    def similarity_search_with_score_by_vector(
        self,
        embedding: Optional[List[float]],
        k: int = 4,
        filter: Optional[Union[dict, list]] = None,
        predicates: Optional[Predicates] = None,
        **kwargs: Any,
    ) -> List[Tuple[Document, float]]:
        try:
            from timescale_vector import client
        except ImportError:
            raise ImportError(
                "Could not import timescale_vector python package. "
                "Please install it with `pip install timescale-vector`."
            )

        results = self.sync_client.search(
            embedding,
            limit=k,
            filter=filter,
            predicates=predicates,
            uuid_time_filter=self.date_to_range_filter(**kwargs),
        )

        docs = [
            (
                Document(
                    page_content=result[client.SEARCH_RESULT_CONTENTS_IDX],
                    metadata=result[client.SEARCH_RESULT_METADATA_IDX],
                ),
                result[client.SEARCH_RESULT_DISTANCE_IDX],
            )
            for result in results
        ]
        return docs

    async def asimilarity_search_with_score_by_vector(
        self,
        embedding: Optional[List[float]],
        k: int = 4,
        filter: Optional[Union[dict, list]] = None,
        predicates: Optional[Predicates] = None,
        **kwargs: Any,
    ) -> List[Tuple[Document, float]]:
        try:
            from timescale_vector import client
        except ImportError:
            raise ImportError(
                "Could not import timescale_vector python package. "
                "Please install it with `pip install timescale-vector`."
            )

        results = await self.async_client.search(
            embedding,
            limit=k,
            filter=filter,
            predicates=predicates,
            uuid_time_filter=self.date_to_range_filter(**kwargs),
        )

        docs = [
            (
                Document(
                    page_content=result[client.SEARCH_RESULT_CONTENTS_IDX],
                    metadata=result[client.SEARCH_RESULT_METADATA_IDX],
                ),
                result[client.SEARCH_RESULT_DISTANCE_IDX],
            )
            for result in results
        ]
        return docs

    def similarity_search_by_vector(
        self,
        embedding: Optional[List[float]],
        k: int = 4,
        filter: Optional[Union[dict, list]] = None,
        predicates: Optional[Predicates] = None,
        **kwargs: Any,
    ) -> List[Document]:
        """Return docs most similar to embedding vector.

        Args:
            embedding: Embedding to look up documents similar to.
            k: Number of Documents to return. Defaults to 4.
            filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.

        Returns:
            List of Documents most similar to the query vector.
        """
        docs_and_scores = self.similarity_search_with_score_by_vector(
            embedding=embedding, k=k, filter=filter, predicates=predicates, **kwargs
        )
        return [doc for doc, _ in docs_and_scores]

    async def asimilarity_search_by_vector(
        self,
        embedding: Optional[List[float]],
        k: int = 4,
        filter: Optional[Union[dict, list]] = None,
        predicates: Optional[Predicates] = None,
        **kwargs: Any,
    ) -> List[Document]:
        """Return docs most similar to embedding vector.

        Args:
            embedding: Embedding to look up documents similar to.
            k: Number of Documents to return. Defaults to 4.
            filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None.

        Returns:
            List of Documents most similar to the query vector.
        """
        docs_and_scores = await self.asimilarity_search_with_score_by_vector(
            embedding=embedding, k=k, filter=filter, predicates=predicates, **kwargs
        )
        return [doc for doc, _ in docs_and_scores]

    @classmethod
    def from_texts(
        cls: Type[TimescaleVector],
        texts: List[str],
        embedding: Embeddings,
        metadatas: Optional[List[dict]] = None,
        collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME,
        distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY,
        ids: Optional[List[str]] = None,
        pre_delete_collection: bool = False,
        **kwargs: Any,
    ) -> TimescaleVector:
        """
        Return VectorStore initialized from texts and embeddings.
        Postgres connection string is required
        "Either pass it as a parameter
        or set the TIMESCALE_SERVICE_URL environment variable.
        """
        embeddings = embedding.embed_documents(list(texts))

        return cls.__from(
            texts,
            embeddings,
            embedding,
            metadatas=metadatas,
            ids=ids,
            collection_name=collection_name,
            distance_strategy=distance_strategy,
            pre_delete_collection=pre_delete_collection,
            **kwargs,
        )

    @classmethod
    async def afrom_texts(
        cls: Type[TimescaleVector],
        texts: List[str],
        embedding: Embeddings,
        metadatas: Optional[List[dict]] = None,
        collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME,
        distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY,
        ids: Optional[List[str]] = None,
        pre_delete_collection: bool = False,
        **kwargs: Any,
    ) -> TimescaleVector:
        """
        Return VectorStore initialized from texts and embeddings.
        Postgres connection string is required
        "Either pass it as a parameter
        or set the TIMESCALE_SERVICE_URL environment variable.
        """
        embeddings = embedding.embed_documents(list(texts))

        return await cls.__afrom(
            texts,
            embeddings,
            embedding,
            metadatas=metadatas,
            ids=ids,
            collection_name=collection_name,
            distance_strategy=distance_strategy,
            pre_delete_collection=pre_delete_collection,
            **kwargs,
        )

    @classmethod
    def from_embeddings(
        cls,
        text_embeddings: List[Tuple[str, List[float]]],
        embedding: Embeddings,
        metadatas: Optional[List[dict]] = None,
        collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME,
        distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY,
        ids: Optional[List[str]] = None,
        pre_delete_collection: bool = False,
        **kwargs: Any,
    ) -> TimescaleVector:
        """Construct TimescaleVector wrapper from raw documents and pre-
        generated embeddings.

        Return VectorStore initialized from documents and embeddings.
        Postgres connection string is required
        "Either pass it as a parameter
        or set the TIMESCALE_SERVICE_URL environment variable.

        Example:
            .. code-block:: python

                from langchain_community.vectorstores import TimescaleVector
                from langchain_community.embeddings import OpenAIEmbeddings
                embeddings = OpenAIEmbeddings()
                text_embeddings = embeddings.embed_documents(texts)
                text_embedding_pairs = list(zip(texts, text_embeddings))
                tvs = TimescaleVector.from_embeddings(text_embedding_pairs, embeddings)
        """
        texts = [t[0] for t in text_embeddings]
        embeddings = [t[1] for t in text_embeddings]

        return cls.__from(
            texts,
            embeddings,
            embedding,
            metadatas=metadatas,
            ids=ids,
            collection_name=collection_name,
            distance_strategy=distance_strategy,
            pre_delete_collection=pre_delete_collection,
            **kwargs,
        )

    @classmethod
    async def afrom_embeddings(
        cls,
        text_embeddings: List[Tuple[str, List[float]]],
        embedding: Embeddings,
        metadatas: Optional[List[dict]] = None,
        collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME,
        distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY,
        ids: Optional[List[str]] = None,
        pre_delete_collection: bool = False,
        **kwargs: Any,
    ) -> TimescaleVector:
        """Construct TimescaleVector wrapper from raw documents and pre-
        generated embeddings.

        Return VectorStore initialized from documents and embeddings.
        Postgres connection string is required
        "Either pass it as a parameter
        or set the TIMESCALE_SERVICE_URL environment variable.

        Example:
            .. code-block:: python

                from langchain_community.vectorstores import TimescaleVector
                from langchain_community.embeddings import OpenAIEmbeddings
                embeddings = OpenAIEmbeddings()
                text_embeddings = embeddings.embed_documents(texts)
                text_embedding_pairs = list(zip(texts, text_embeddings))
                tvs = TimescaleVector.from_embeddings(text_embedding_pairs, embeddings)
        """
        texts = [t[0] for t in text_embeddings]
        embeddings = [t[1] for t in text_embeddings]

        return await cls.__afrom(
            texts,
            embeddings,
            embedding,
            metadatas=metadatas,
            ids=ids,
            collection_name=collection_name,
            distance_strategy=distance_strategy,
            pre_delete_collection=pre_delete_collection,
            **kwargs,
        )

    @classmethod
    def from_existing_index(
        cls: Type[TimescaleVector],
        embedding: Embeddings,
        collection_name: str = _LANGCHAIN_DEFAULT_COLLECTION_NAME,
        distance_strategy: DistanceStrategy = DEFAULT_DISTANCE_STRATEGY,
        pre_delete_collection: bool = False,
        **kwargs: Any,
    ) -> TimescaleVector:
        """
        Get instance of an existing TimescaleVector store.This method will
        return the instance of the store without inserting any new
        embeddings
        """

        service_url = cls.get_service_url(kwargs)

        store = cls(
            service_url=service_url,
            collection_name=collection_name,
            embedding=embedding,
            distance_strategy=distance_strategy,
            pre_delete_collection=pre_delete_collection,
        )

        return store

    @classmethod
    def get_service_url(cls, kwargs: Dict[str, Any]) -> str:
        service_url: str = get_from_dict_or_env(
            data=kwargs,
            key="service_url",
            env_key="TIMESCALE_SERVICE_URL",
        )

        if not service_url:
            raise ValueError(
                "Postgres connection string is required"
                "Either pass it as a parameter"
                "or set the TIMESCALE_SERVICE_URL environment variable."
            )

        return service_url

    @classmethod
    def service_url_from_db_params(
        cls,
        host: str,
        port: int,
        database: str,
        user: str,
        password: str,
    ) -> str:
        """Return connection string from database parameters."""
        return f"postgresql://{user}:{password}@{host}:{port}/{database}"

    def _select_relevance_score_fn(self) -> Callable[[float], float]:
        """
        The 'correct' relevance function
        may differ depending on a few things, including:
        - the distance / similarity metric used by the VectorStore
        - the scale of your embeddings (OpenAI's are unit normed. Many others are not!)
        - embedding dimensionality
        - etc.
        """
        if self.override_relevance_score_fn is not None:
            return self.override_relevance_score_fn

        # Default strategy is to rely on distance strategy provided
        # in vectorstore constructor
        if self._distance_strategy == DistanceStrategy.COSINE:
            return self._cosine_relevance_score_fn
        elif self._distance_strategy == DistanceStrategy.EUCLIDEAN_DISTANCE:
            return self._euclidean_relevance_score_fn
        elif self._distance_strategy == DistanceStrategy.MAX_INNER_PRODUCT:
            return self._max_inner_product_relevance_score_fn
        else:
            raise ValueError(
                "No supported normalization function"
                f" for distance_strategy of {self._distance_strategy}."
                "Consider providing relevance_score_fn to TimescaleVector constructor."
            )

    def delete(self, ids: Optional[List[str]] = None, **kwargs: Any) -> Optional[bool]:
        """Delete by vector ID or other criteria.

        Args:
            ids: List of ids to delete.
            **kwargs: Other keyword arguments that subclasses might use.

        Returns:
            Optional[bool]: True if deletion is successful,
            False otherwise, None if not implemented.
        """
        if ids is None:
            raise ValueError("No ids provided to delete.")

        self.sync_client.delete_by_ids(ids)
        return True

    # todo should this be part of delete|()?
    def delete_by_metadata(
        self, filter: Union[Dict[str, str], List[Dict[str, str]]], **kwargs: Any
    ) -> Optional[bool]:
        """Delete by vector ID or other criteria.

        Args:
            ids: List of ids to delete.
            **kwargs: Other keyword arguments that subclasses might use.

        Returns:
            Optional[bool]: True if deletion is successful,
            False otherwise, None if not implemented.
        """

        self.sync_client.delete_by_metadata(filter)
        return True

    class IndexType(str, enum.Enum):
        """Enumerator for the supported Index types"""

        TIMESCALE_VECTOR = "tsv"
        PGVECTOR_IVFFLAT = "ivfflat"
        PGVECTOR_HNSW = "hnsw"

    DEFAULT_INDEX_TYPE = IndexType.TIMESCALE_VECTOR

    def create_index(
        self, index_type: Union[IndexType, str] = DEFAULT_INDEX_TYPE, **kwargs: Any
    ) -> None:
        try:
            from timescale_vector import client
        except ImportError:
            raise ImportError(
                "Could not import timescale_vector python package. "
                "Please install it with `pip install timescale-vector`."
            )

        index_type = (
            index_type.value if isinstance(index_type, self.IndexType) else index_type
        )
        if index_type == self.IndexType.PGVECTOR_IVFFLAT.value:
            self.sync_client.create_embedding_index(client.IvfflatIndex(**kwargs))

        if index_type == self.IndexType.PGVECTOR_HNSW.value:
            self.sync_client.create_embedding_index(client.HNSWIndex(**kwargs))

        if index_type == self.IndexType.TIMESCALE_VECTOR.value:
            self.sync_client.create_embedding_index(
                client.TimescaleVectorIndex(**kwargs)
            )

    def drop_index(self) -> None:
        self.sync_client.drop_embedding_index()