Manages per-document vector-store collections. Each uploaded PDF gets its own ChromaDB collection, keyed by a document ID (typically an MD5 hash). Collections can live in memory or be persisted to disk. Args: embedding_function: A LangChain-compatible ``Embeddings`` instan
| 158 | |
| 159 | |
| 160 | class DataStorage: |
| 161 | """Manages per-document vector-store collections. |
| 162 | |
| 163 | Each uploaded PDF gets its own ChromaDB collection, |
| 164 | keyed by a document ID (typically an MD5 hash). Collections can live |
| 165 | in memory or be persisted to disk. |
| 166 | |
| 167 | Args: |
| 168 | embedding_function: A LangChain-compatible ``Embeddings`` instance |
| 169 | root_path: Optional directory for persisted embeddings. |
| 170 | engine: The vector-store class to use. |
| 171 | |
| 172 | """ |
| 173 | |
| 174 | embeddings_dict = {} |
| 175 | embeddings_map_from_md5 = {} |
| 176 | embeddings_map_to_md5 = {} |
| 177 | |
| 178 | def __init__( |
| 179 | self, |
| 180 | embedding_function, |
| 181 | root_path: Path = None, |
| 182 | engine=ChromaAdvancedRetrieval, |
| 183 | ) -> None: |
| 184 | self.root_path = root_path |
| 185 | self.engine = engine |
| 186 | self.embedding_function = embedding_function |
| 187 | |
| 188 | if root_path is not None: |
| 189 | self.embeddings_root_path = root_path |
| 190 | if not os.path.exists(root_path): |
| 191 | os.makedirs(root_path) |
| 192 | else: |
| 193 | self.load_embeddings(self.embeddings_root_path) |
| 194 | |
| 195 | def load_embeddings(self, embeddings_root_path: Union[str, Path]) -> None: |
| 196 | """ |
| 197 | Load the vector storage assuming they are all persisted and stored in a single directory. |
| 198 | The root path of the embeddings containing one data store for each document in each subdirectory |
| 199 | """ |
| 200 | |
| 201 | embeddings_directories = [f for f in os.scandir(embeddings_root_path) if f.is_dir()] |
| 202 | |
| 203 | if len(embeddings_directories) == 0: |
| 204 | print("No available embeddings") |
| 205 | return |
| 206 | |
| 207 | for embedding_document_dir in embeddings_directories: |
| 208 | self.embeddings_dict[embedding_document_dir.name] = self.engine( |
| 209 | persist_directory=embedding_document_dir.path, embedding_function=self.embedding_function |
| 210 | ) |
| 211 | |
| 212 | filename_list = list(Path(embedding_document_dir).glob("*.storage_filename")) |
| 213 | if filename_list: |
| 214 | filenam = filename_list[0].name.replace(".storage_filename", "") |
| 215 | self.embeddings_map_from_md5[embedding_document_dir.name] = filenam |
| 216 | self.embeddings_map_to_md5[filenam] = embedding_document_dir.name |
| 217 |