Repository: SuperDuperDB/superduperdb
File: superduperdb/vector_search/atlas.py (Python)

Code before:

if t.TYPE_CHECKING:
    from superduperdb.components.vector_index import VectorIndex


class MongoAtlasVectorSearcher(BaseVectorSearcher):
    """
    Implementation of atlas vector search

    :param identifier: Unique string identifier of index
    """

    def __init__(
        self,
        identifier: str,
        collection: str,
        dimensions: t.Optional[int] = None,
        measure: t.Optional[str] = None,
        output_path: t.Optional[str] = None,
        **kwargs,
    ):
        self.identifier = identifier
        db_name = CFG.cluster.vector_search.split('/')[-1]
        self.database = getattr(pymongo.MongoClient(CFG.cluster.vector_search), db_name)
        assert output_path
        self.output_path = output_path
        self.collection = collection
        self.measure = measure
        self.dimensions = dimensions

        if not self._check_if_exists(identifier):
            self._create_index(collection, output_path)

    def __len__(self):
        pass

    @cached_property
    def index(self):
        return self.database[self.collection]

    @classmethod
    def from_component(cls, vi: 'VectorIndex'):
        from superduperdb.components.listener import Listener
        from superduperdb.components.model import ObjectModel

        assert isinstance(vi.indexing_listener, Listener)
        collection = vi.indexing_listener.select.table_or_collection.identifier

        indexing_key = vi.indexing_listener.key
        assert isinstance(
            indexing_key, str

Review comment (suggestion):
After I added the type reminder for Config `CFG: Config = ConfigSettings(Config).config`, an error occurred during detection here. It may be because we changed the structure of config and did not update it here. I fixed it, but did not test it.

Code after:

if t.TYPE_CHECKING:
    from superduperdb.components.vector_index import VectorIndex


class MongoAtlasVectorSearcher(BaseVectorSearcher):
    """
    Implementation of atlas vector search

    :param identifier: Unique string identifier of index
    """

    def __init__(
        self,
        identifier: str,
        collection: str,
        dimensions: t.Optional[int] = None,
        measure: t.Optional[str] = None,
        output_path: t.Optional[str] = None,
        **kwargs,
    ):
        self.identifier = identifier
        vector_search_uri = CFG.cluster.vector_search.uri
        assert vector_search_uri, 'Vector search URI is required'
        db_name = vector_search_uri.split('/')[-1]
        self.database = getattr(pymongo.MongoClient(vector_search_uri), db_name)
        assert output_path
        self.output_path = output_path
        self.collection = collection
        self.measure = measure
        self.dimensions = dimensions

        if not self._check_if_exists(identifier):
            self._create_index(collection, output_path)

    def __len__(self):
        pass

    @cached_property
    def index(self):
        return self.database[self.collection]

    @classmethod
    def from_component(cls, vi: 'VectorIndex'):
        from superduperdb.components.listener import Listener
        from superduperdb.components.model import ObjectModel

        assert isinstance(vi.indexing_listener, Listener)
        collection = vi.indexing_listener.select.table_or_collection.identifier

        indexing_key = vi.indexing_listener.key