llama-index

1
from enum import Enum
2
from typing import Dict, Type
3

4
from llama_index.legacy.storage.docstore.mongo_docstore import MongoDocumentStore
5
from llama_index.legacy.storage.docstore.simple_docstore import SimpleDocumentStore
6
from llama_index.legacy.storage.docstore.types import BaseDocumentStore
7

8

9
class DocumentStoreType(str, Enum):
10
    MONGO = "mongo"
11
    SIMPLE = "simple"
12

13

14
DOCSTORE_TYPE_TO_CLASS: Dict[DocumentStoreType, Type[BaseDocumentStore]] = {
15
    DocumentStoreType.MONGO: MongoDocumentStore,
16
    DocumentStoreType.SIMPLE: SimpleDocumentStore,
17
}
18

19

20
DOCSTORE_CLASS_TO_TYPE: Dict[Type[BaseDocumentStore], DocumentStoreType] = {
21
    cls_: type_ for type_, cls_ in DOCSTORE_TYPE_TO_CLASS.items()
22
}
23

24

25
def get_default_docstore() -> BaseDocumentStore:
26
    return SimpleDocumentStore()
27