-
Notifications
You must be signed in to change notification settings - Fork 315
/
SimpleKVStore.ts
109 lines (95 loc) · 2.81 KB
/
SimpleKVStore.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
import type { GenericFileSystem } from "@llamaindex/env";
import { defaultFS, path } from "@llamaindex/env";
import _ from "lodash";
import { exists } from "../FileSystem.js";
import { DEFAULT_COLLECTION } from "../constants.js";
import { BaseKVStore } from "./types.js";
export type DataType = Record<string, Record<string, any>>;
export class SimpleKVStore extends BaseKVStore {
private persistPath: string | undefined;
private fs: GenericFileSystem | undefined;
constructor(private data: DataType = {}) {
super();
}
async put(
key: string,
val: any,
collection: string = DEFAULT_COLLECTION,
): Promise<void> {
if (!(collection in this.data)) {
this.data[collection] = {};
}
this.data[collection][key] = _.clone(val); // Creating a shallow copy of the object
if (this.persistPath) {
await this.persist(this.persistPath, this.fs);
}
}
async get(
key: string,
collection: string = DEFAULT_COLLECTION,
): Promise<any> {
const collectionData = this.data[collection];
if (_.isNil(collectionData)) {
return null;
}
if (!(key in collectionData)) {
return null;
}
return _.clone(collectionData[key]); // Creating a shallow copy of the object
}
async getAll(collection: string = DEFAULT_COLLECTION): Promise<DataType> {
return _.clone(this.data[collection]); // Creating a shallow copy of the object
}
async delete(
key: string,
collection: string = DEFAULT_COLLECTION,
): Promise<boolean> {
if (key in this.data[collection]) {
delete this.data[collection][key];
if (this.persistPath) {
await this.persist(this.persistPath, this.fs);
}
return true;
}
return false;
}
async persist(
persistPath: string,
fs: GenericFileSystem = defaultFS,
): Promise<void> {
// TODO: decide on a way to polyfill path
const dirPath = path.dirname(persistPath);
if (!(await exists(fs, dirPath))) {
await fs.mkdir(dirPath);
}
await fs.writeFile(persistPath, JSON.stringify(this.data));
}
static async fromPersistPath(
persistPath: string,
fs: GenericFileSystem = defaultFS,
): Promise<SimpleKVStore> {
const dirPath = path.dirname(persistPath);
if (!(await exists(fs, dirPath))) {
await fs.mkdir(dirPath);
}
let data: DataType = {};
try {
const fileData = await fs.readFile(persistPath);
data = JSON.parse(fileData.toString());
} catch (e) {
console.error(
`No valid data found at path: ${persistPath} starting new store.`,
);
}
const store = new SimpleKVStore(data);
store.persistPath = persistPath;
store.fs = fs;
return store;
}
toDict(): DataType {
return this.data;
}
static fromDict(saveDict: DataType): SimpleKVStore {
return new SimpleKVStore(saveDict);
}
}