-
Notifications
You must be signed in to change notification settings - Fork 311
/
json-to-index-struct.ts
78 lines (68 loc) · 1.96 KB
/
json-to-index-struct.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
import type { BaseNode } from "../Node.js";
import { jsonToNode } from "../Node.js";
import { IndexStruct } from "./IndexStruct.js";
export enum IndexStructType {
SIMPLE_DICT = "simple_dict",
LIST = "list",
KEYWORD_TABLE = "keyword_table",
}
export class IndexDict extends IndexStruct {
nodesDict: Record<string, BaseNode> = {};
type: IndexStructType = IndexStructType.SIMPLE_DICT;
getSummary(): string {
if (this.summary === undefined) {
throw new Error("summary field of the index dict is not set");
}
return this.summary;
}
addNode(node: BaseNode, textId?: string) {
const vectorId = textId ?? node.id_;
this.nodesDict[vectorId] = node;
}
toJson(): Record<string, unknown> {
const nodesDict: Record<string, unknown> = {};
for (const [key, node] of Object.entries(this.nodesDict)) {
nodesDict[key] = node.toJSON();
}
return {
...super.toJson(),
nodesDict,
type: this.type,
};
}
delete(nodeId: string) {
delete this.nodesDict[nodeId];
}
}
export function jsonToIndexStruct(json: any): IndexStruct {
if (json.type === IndexStructType.LIST) {
const indexList = new IndexList(json.indexId, json.summary);
indexList.nodes = json.nodes;
return indexList;
} else if (json.type === IndexStructType.SIMPLE_DICT) {
const indexDict = new IndexDict(json.indexId, json.summary);
indexDict.nodesDict = Object.entries(json.nodesDict).reduce<
Record<string, BaseNode>
>((acc, [key, value]) => {
acc[key] = jsonToNode(value);
return acc;
}, {});
return indexDict;
} else {
throw new Error(`Unknown index struct type: ${json.type}`);
}
}
export class IndexList extends IndexStruct {
nodes: string[] = [];
type: IndexStructType = IndexStructType.LIST;
addNode(node: BaseNode) {
this.nodes.push(node.id_);
}
toJson(): Record<string, unknown> {
return {
...super.toJson(),
nodes: this.nodes,
type: this.type,
};
}
}