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

Code before:
        -----------

            - collection_name (str): The name of the new collection to create.
            - payload_schema (Optional[Any]): An optional schema for the payloads
              within this collection. Can be PayloadSchema, BaseModel, or other schema types. (default None)
        """
        raise NotImplementedError

    """ Data points """

    @abstractmethod
    async def create_data_points(self, collection_name: str, data_points: List[DataPoint]):
        """
        Insert new data points into the specified collection.

        Parameters:
        -----------

            - collection_name (str): The name of the collection where data points will be added.
            - data_points (List[DataPoint]): A list of data points to be added to the
              collection.
        """
        raise NotImplementedError

    @abstractmethod
    async def retrieve(self, collection_name: str, data_point_ids: Union[List[str], list[str]]):
        """
        Retrieve data points from a collection using their IDs.

        Parameters:
        -----------

            - collection_name (str): The name of the collection from which to retrieve data
              points.
            - data_point_ids (Union[List[str], list[str]]): A list of IDs of the data points to retrieve.
        """
        raise NotImplementedError

    """ Search """

    @abstractmethod
    async def search(
        self,
        collection_name: str,
        query_text: Optional[str],
        query_vector: Optional[List[float]],
        limit: int,
        with_vector: bool = False,
    ):
        """
        Perform a search in the specified collection using either a text query or a vector

Review comment (suggestion):
```suggestion
    async def retrieve(self, collection_name: str, data_point_ids: list[str]):
```

Noticed since we're on 3.9+ Python can keep just this one to keep it simpler

Code after:
        -----------

            - collection_name (str): The name of the new collection to create.
            - payload_schema (Optional[Any]): An optional schema for the payloads
              within this collection. Can be PayloadSchema, BaseModel, or other schema types. (default None)
        """
        raise NotImplementedError

    """ Data points """

    @abstractmethod
    async def create_data_points(self, collection_name: str, data_points: List[DataPoint]):
        """
        Insert new data points into the specified collection.

        Parameters:
        -----------

            - collection_name (str): The name of the collection where data points will be added.
            - data_points (List[DataPoint]): A list of data points to be added to the
              collection.
        """
        raise NotImplementedError

    @abstractmethod
    async def retrieve(self, collection_name: str, data_point_ids: list[str]):
        """
        Retrieve data points from a collection using their IDs.

        Parameters:
        -----------

            - collection_name (str): The name of the collection from which to retrieve data
              points.
            - data_point_ids (Union[List[str], list[str]]): A list of IDs of the data points to retrieve.
        """
        raise NotImplementedError

    """ Search """

    @abstractmethod
    async def search(
        self,
        collection_name: str,
        query_text: Optional[str],
        query_vector: Optional[List[float]],
        limit: int,
        with_vector: bool = False,
    ):
        """
        Perform a search in the specified collection using either a text query or a vector