diff --git a/.changeset/eight-dogs-burn.md b/.changeset/eight-dogs-burn.md new file mode 100644 index 0000000000000000000000000000000000000000..71435c76bbd58ee71058f101a217555e7e036f7b --- /dev/null +++ b/.changeset/eight-dogs-burn.md @@ -0,0 +1,12 @@ +--- +"llamaindex": patch +"@llamaindex/core-test": patch +--- + +- Add missing exports: + - `IndexStructType`, + - `IndexDict`, + - `jsonToIndexStruct`, + - `IndexList`, + - `IndexStruct` +- Fix `IndexDict.toJson()` method diff --git a/packages/core/src/indices/index.ts b/packages/core/src/indices/index.ts index 9fce474e42c7d132cb9634d4559fd71f6735914e..196287e2ab5991a44ddc67321c625a074013d7f1 100644 --- a/packages/core/src/indices/index.ts +++ b/packages/core/src/indices/index.ts @@ -1,4 +1,6 @@ export * from "./BaseIndex.js"; +export * from "./IndexStruct.js"; +export * from "./json-to-index-struct.js"; export * from "./keyword/index.js"; export * from "./summary/index.js"; export * from "./vectorStore/index.js"; diff --git a/packages/core/src/indices/json-to-index-struct.ts b/packages/core/src/indices/json-to-index-struct.ts index b62ec851bd8149bf38c47dd099a550098b9735f0..158c3af73e47109f1ba72c6e30ba5b18b4c4890f 100644 --- a/packages/core/src/indices/json-to-index-struct.ts +++ b/packages/core/src/indices/json-to-index-struct.ts @@ -24,9 +24,15 @@ export class IndexDict extends IndexStruct { } 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: this.nodesDict, + nodesDict, type: this.type, }; } diff --git a/packages/core/tests/indices/json-to-index-struct.test.ts b/packages/core/tests/indices/json-to-index-struct.test.ts new file mode 100644 index 0000000000000000000000000000000000000000..ed996e9790d57b92b13c6857c35d913263f9a12b --- /dev/null +++ b/packages/core/tests/indices/json-to-index-struct.test.ts @@ -0,0 +1,71 @@ +import { + IndexDict, + IndexList, + IndexStruct, + IndexStructType, + MetadataMode, + TextNode, + jsonToIndexStruct, +} from "llamaindex"; +import { describe, expect, it } from "vitest"; + +describe("jsonToIndexStruct", () => { + it("transforms json to IndexDict", () => { + function isIndexDict(some: IndexStruct): some is IndexDict { + return "type" in some && some.type === IndexStructType.SIMPLE_DICT; + } + + const node = new TextNode({ text: "text", id_: "nodeId" }); + const expected = new IndexDict(); + expected.addNode(node); + + console.log("expected.toJson()", expected.toJson()); + const actual = jsonToIndexStruct(expected.toJson()); + + expect(isIndexDict(actual)).toBe(true); + expect( + (actual as IndexDict).nodesDict.nodeId.getContent(MetadataMode.NONE), + ).toEqual("text"); + }); + it("transforms json to IndexList", () => { + function isIndexList(some: IndexStruct): some is IndexList { + return "type" in some && some.type === IndexStructType.LIST; + } + + const node = new TextNode({ text: "text", id_: "nodeId" }); + const expected = new IndexList(); + expected.addNode(node); + + const actual = jsonToIndexStruct(expected.toJson()); + + expect(isIndexList(actual)).toBe(true); + expect((actual as IndexList).nodes[0]).toEqual("nodeId"); + }); + it("fails for unknown index type", () => { + expect(() => { + const json = { + indexId: "dd120b16-8dce-4ce3-9bb6-15ca87fe4a1d", + summary: undefined, + nodesDict: {}, + type: "FOO", + }; + return jsonToIndexStruct(json); + }).toThrowError("Unknown index struct type: FOO"); + }); + it("fails for unknown node type", () => { + expect(() => { + const json = { + indexId: "dd120b16-8dce-4ce3-9bb6-15ca87fe4a1d", + summary: undefined, + nodesDict: { + nodeId: { + ...new TextNode({ text: "text", id_: "nodeId" }).toJSON(), + type: "BAR", + }, + }, + type: IndexStructType.SIMPLE_DICT, + }; + return jsonToIndexStruct(json); + }).toThrowError("Invalid node type: BAR"); + }); +});