Skip to content

momento

langroid/vector_store/momento.py

Momento Vector Index. https://docs.momentohq.com/vector-index/develop/api-reference

MomentoVI(config=MomentoVIConfig())

Bases: VectorStore

Source code in langroid/vector_store/momento.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
def __init__(self, config: MomentoVIConfig = MomentoVIConfig()):
    super().__init__(config)
    self.config: MomentoVIConfig = config
    emb_model = EmbeddingModel.create(config.embedding)
    self.embedding_fn: EmbeddingFunction = emb_model.embedding_fn()
    self.embedding_dim = emb_model.embedding_dims
    self.host = config.host
    self.port = config.port
    load_dotenv()
    api_key = os.getenv("MOMENTO_API_KEY")
    if config.cloud:
        if api_key is None:
            raise ValueError(
                """MOMENTO_API_KEY env variable must be set to 
                MomentoVI hosted service. Please set this in your .env file. 
                """
            )
        self.client = PreviewVectorIndexClient(
            configuration=VectorIndexConfigurations.Default.latest(),
            credential_provider=CredentialProvider.from_string(api_key),
        )
    else:
        raise NotImplementedError("MomentoVI local not available yet")

    # Note: Only create collection if a non-null collection name is provided.
    # This is useful to delay creation of vecdb until we have a suitable
    # collection name (e.g. we could get it from the url or folder path).
    if config.collection_name is not None:
        self.create_collection(
            config.collection_name, replace=config.replace_collection
        )

clear_all_collections(really=False, prefix='')

Clear all collections with the given prefix.

Source code in langroid/vector_store/momento.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
def clear_all_collections(self, really: bool = False, prefix: str = "") -> int:
    """Clear all collections with the given prefix."""

    if not really:
        logger.warning("Not deleting all collections, set really=True to confirm")
        return 0
    coll_names = self.list_collections(empty=False)
    coll_names = [name for name in coll_names if name.startswith(prefix)]
    if len(coll_names) == 0:
        logger.warning(f"No collections found with prefix {prefix}")
        return 0
    for name in coll_names:
        self.delete_collection(name)
    logger.warning(
        f"""
        Deleted {len(coll_names)} indices from Momento VI
        """
    )
    return len(coll_names)

list_collections(empty=False)

Returns:

Type Description
List[str]

List of collection names that have at least one vector.

Parameters:

Name Type Description Default
empty bool

Whether to include empty collections.

 False
Source code in langroid/vector_store/momento.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
def list_collections(self, empty: bool = False) -> List[str]:
    """
    Returns:
        List of collection names that have at least one vector.

    Args:
        empty (bool, optional): Whether to include empty collections.
    """
    response = self.client.list_indexes()
    if isinstance(response, mvi_response.ListIndexes.Success):
        return [ind.name for ind in response.indexes]
    elif isinstance(response, mvi_response.ListIndexes.Error):
        raise ValueError(f"Error listing collections: {response.message}")
    else:
        raise ValueError(f"Unexpected response: {response}")

create_collection(collection_name, replace=False)

Create a collection with the given name, optionally replacing an existing collection if replace is True. Args: collection_name (str): Name of the collection to create. replace (bool): Whether to replace an existing collection with the same name. Defaults to False.

Source code in langroid/vector_store/momento.py 
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
def create_collection(self, collection_name: str, replace: bool = False) -> None:
    """
    Create a collection with the given name, optionally replacing an existing
        collection if `replace` is True.
    Args:
        collection_name (str): Name of the collection to create.
        replace (bool): Whether to replace an existing collection
            with the same name. Defaults to False.
    """
    self.config.collection_name = collection_name
    response = self.client.create_index(
        index_name=collection_name,
        num_dimensions=self.embedding_dim,
        similarity_metric=self.config.distance,
    )
    if isinstance(response, mvi_response.CreateIndex.Success):
        logger.info(f"Created collection {collection_name}")
    elif isinstance(response, mvi_response.CreateIndex.IndexAlreadyExists):
        logger.warning(f"Collection {collection_name} already exists")
    elif isinstance(response, mvi_response.CreateIndex.Error):
        raise ValueError(
            f"Error creating collection {collection_name}: {response.message}"
        )
    if settings.debug:
        level = logger.getEffectiveLevel()
        logger.setLevel(logging.INFO)
        logger.info(f"Collection {collection_name} created")
        logger.setLevel(level)