fix: legacy js object map

run-llama · Sep 4, 2024 · dc77d97 · dc77d97
1 parent 68d0f71
commit dc77d97
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 30 deletions.
diff --git a/packages/core/src/schema/index.ts b/packages/core/src/schema/index.ts
@@ -1,5 +1,5 @@
 export * from "./node";
-export { DATA_KEY, TYPE_KEY, fromPythonDocStore, type DocJson } from "./python";
+export { DATA_KEY, TYPE_KEY, fromDocStore, type DocJson } from "./python";
 export { FileReader, TransformComponent, type BaseReader } from "./type";
 export { EngineResponse } from "./type/engine–response";
 export * from "./zod";
diff --git a/packages/core/src/schema/python.ts b/packages/core/src/schema/python.ts
@@ -28,12 +28,20 @@ const PYTHON_TO_JS_TYPE_MAP = {
   "4": ObjectType.DOCUMENT,
 };
 
+const LEGACY_JS_MAP = {
+  TEXT: ObjectType.TEXT,
+  IMAGE: ObjectType.IMAGE,
+  INDEX: ObjectType.INDEX,
+  DOCUMENT: ObjectType.DOCUMENT,
+  IMAGE_DOCUMENT: ObjectType.DOCUMENT,
+};
+
 export type DocJson = {
   [TYPE_KEY]: string;
   [DATA_KEY]: string;
 };
 
-async function fromPythonImpl(data: Record<string, unknown>) {
+async function fromImpl(data: Record<string, unknown>) {
   const convertedJson = await camelCaseJson(data);
   if (convertedJson.relationships) {
     for (const [key, value] of Object.entries(convertedJson.relationships)) {
@@ -49,26 +57,16 @@ async function fromPythonImpl(data: Record<string, unknown>) {
   return convertedJson;
 }
 
-export async function fromPythonDocStore({
+export async function fromDocStore({
   [TYPE_KEY]: type,
   [DATA_KEY]: data,
 }: DocJson) {
-  if (!(type in PYTHON_TO_JS_TYPE_MAP)) {
-    throw new Error("");
-  }
-  const objectType =
-    PYTHON_TO_JS_TYPE_MAP[type as keyof typeof PYTHON_TO_JS_TYPE_MAP];
-  const convertedJson = await fromPythonImpl(JSON.parse(data));
-  return jsonToNode(convertedJson, objectType);
-}
-
-export async function fromPythonNode(json: Record<any, unknown>) {
-  const convertedJson = await fromPythonImpl(json);
-  const type = convertedJson["metadata"]?.["_node_type"];
-  if (!(type in PYTHON_TO_JS_TYPE_MAP)) {
-    throw new Error("");
+  if (!(type in PYTHON_TO_JS_TYPE_MAP) && !(type in LEGACY_JS_MAP)) {
+    throw new Error("Invalid type");
   }
   const objectType =
-    PYTHON_TO_JS_TYPE_MAP[type as keyof typeof PYTHON_TO_JS_TYPE_MAP];
+    PYTHON_TO_JS_TYPE_MAP[type as keyof typeof PYTHON_TO_JS_TYPE_MAP] ||
+    LEGACY_JS_MAP[type as keyof typeof LEGACY_JS_MAP];
+  const convertedJson = await fromImpl(JSON.parse(data));
   return jsonToNode(convertedJson, objectType);
 }
diff --git a/packages/llamaindex/src/ingestion/IngestionCache.ts b/packages/llamaindex/src/ingestion/IngestionCache.ts
@@ -1,7 +1,11 @@
-import type { BaseNode, TransformComponent } from "@llamaindex/core/schema";
-import { MetadataMode } from "@llamaindex/core/schema";
+import {
+  type BaseNode,
+  fromDocStore,
+  MetadataMode,
+  type TransformComponent,
+} from "@llamaindex/core/schema";
 import { createSHA256 } from "@llamaindex/env";
-import { docToJson, jsonToDoc } from "../storage/docStore/utils.js";
+import { docToJson } from "../storage/docStore/utils.js";
 import { SimpleKVStore } from "../storage/kvStore/SimpleKVStore.js";
 import type { BaseKVStore } from "../storage/kvStore/types.js";
 
@@ -63,6 +67,8 @@ export class IngestionCache {
     if (!json || !json[this.nodesKey] || !Array.isArray(json[this.nodesKey])) {
       return undefined;
     }
-    return Promise.all(json[this.nodesKey].map((doc: any) => jsonToDoc(doc)));
+    return Promise.all(
+      json[this.nodesKey].map((doc: any) => fromDocStore(doc)),
+    );
   }
 }
diff --git a/packages/llamaindex/src/storage/docStore/KVDocumentStore.ts b/packages/llamaindex/src/storage/docStore/KVDocumentStore.ts
@@ -1,11 +1,14 @@
-import type { BaseNode } from "@llamaindex/core/schema";
-import { ObjectType } from "@llamaindex/core/schema";
+import {
+  type BaseNode,
+  fromDocStore,
+  ObjectType,
+} from "@llamaindex/core/schema";
 import _ from "lodash";
 import { DEFAULT_NAMESPACE } from "../constants.js";
 import type { BaseKVStore } from "../kvStore/types.js";
 import type { RefDocInfo } from "./types.js";
 import { BaseDocumentStore } from "./types.js";
-import { docToJson, isValidDocJson, jsonToDoc } from "./utils.js";
+import { docToJson, isValidDocJson } from "./utils.js";
 
 type DocMetaData = { docHash: string; refDocId?: string };
 
@@ -29,7 +32,7 @@ export class KVDocumentStore extends BaseDocumentStore {
     for (const key in jsonDict) {
       const value = jsonDict[key];
       if (isValidDocJson(value)) {
-        docs[key] = await jsonToDoc(value);
+        docs[key] = await fromDocStore(value);
       } else {
         console.warn(`Invalid JSON for docId ${key}`);
       }
@@ -94,7 +97,7 @@ export class KVDocumentStore extends BaseDocumentStore {
     if (!isValidDocJson(json)) {
       throw new Error(`Invalid JSON for docId ${docId}`);
     }
-    return jsonToDoc(json);
+    return fromDocStore(json);
   }
 
   async getRefDocInfo(refDocId: string): Promise<RefDocInfo | undefined> {

diff --git a/packages/llamaindex/src/storage/docStore/utils.ts b/packages/llamaindex/src/storage/docStore/utils.ts
@@ -1,5 +1,5 @@
 import type { BaseNode } from "@llamaindex/core/schema";
-import { ObjectType, fromPythonDocStore } from "@llamaindex/core/schema";
+import { ObjectType } from "@llamaindex/core/schema";
 
 const TYPE_KEY = "__type__";
 const DATA_KEY = "__data__";
@@ -24,5 +24,3 @@ export function docToJson(doc: BaseNode): DocJson {
     [TYPE_KEY]: doc.type,
   };
 }
-
-export const jsonToDoc = fromPythonDocStore;