llama-index

1
from typing import List, Optional
2

3
import fsspec
4

5
from llama_index.legacy.bridge.pydantic import BaseModel, Field
6
from llama_index.legacy.schema import BaseNode
7
from llama_index.legacy.storage.docstore.utils import doc_to_json, json_to_doc
8
from llama_index.legacy.storage.kvstore import (
9
    FirestoreKVStore as FirestoreCache,
10
)
11
from llama_index.legacy.storage.kvstore import (
12
    MongoDBKVStore as MongoDBCache,
13
)
14
from llama_index.legacy.storage.kvstore import (
15
    RedisKVStore as RedisCache,
16
)
17
from llama_index.legacy.storage.kvstore import (
18
    SimpleKVStore as SimpleCache,
19
)
20
from llama_index.legacy.storage.kvstore.types import (
21
    BaseKVStore as BaseCache,
22
)
23

24
DEFAULT_CACHE_NAME = "llama_cache"
25

26

27
class IngestionCache(BaseModel):
28
    class Config:
29
        arbitrary_types_allowed = True
30

31
    nodes_key = "nodes"
32

33
    collection: str = Field(
34
        default=DEFAULT_CACHE_NAME, description="Collection name of the cache."
35
    )
36
    cache: BaseCache = Field(default_factory=SimpleCache, description="Cache to use.")
37

38
    # TODO: add async get/put methods?
39
    def put(
40
        self, key: str, nodes: List[BaseNode], collection: Optional[str] = None
41
    ) -> None:
42
        """Put a value into the cache."""
43
        collection = collection or self.collection
44

45
        val = {self.nodes_key: [doc_to_json(node) for node in nodes]}
46
        self.cache.put(key, val, collection=collection)
47

48
    def get(
49
        self, key: str, collection: Optional[str] = None
50
    ) -> Optional[List[BaseNode]]:
51
        """Get a value from the cache."""
52
        collection = collection or self.collection
53
        node_dicts = self.cache.get(key, collection=collection)
54

55
        if node_dicts is None:
56
            return None
57

58
        return [json_to_doc(node_dict) for node_dict in node_dicts[self.nodes_key]]
59

60
    def clear(self, collection: Optional[str] = None) -> None:
61
        """Clear the cache."""
62
        collection = collection or self.collection
63
        data = self.cache.get_all(collection=collection)
64
        for key in data:
65
            self.cache.delete(key, collection=collection)
66

67
    def persist(
68
        self, persist_path: str, fs: Optional[fsspec.AbstractFileSystem] = None
69
    ) -> None:
70
        """Persist the cache to a directory, if possible."""
71
        if isinstance(self.cache, SimpleCache):
72
            self.cache.persist(persist_path, fs=fs)
73
        else:
74
            print("Warning: skipping persist, only needed for SimpleCache.")
75

76
    @classmethod
77
    def from_persist_path(
78
        cls,
79
        persist_path: str,
80
        collection: str = DEFAULT_CACHE_NAME,
81
        fs: Optional[fsspec.AbstractFileSystem] = None,
82
    ) -> "IngestionCache":
83
        """Create a IngestionCache from a persist directory."""
84
        return cls(
85
            collection=collection,
86
            cache=SimpleCache.from_persist_path(persist_path, fs=fs),
87
        )
88

89

90
__all__ = [
91
    "SimpleCache",
92
    "RedisCache",
93
    "MongoDBCache",
94
    "FirestoreCache",
95
]
96