refactor: move storage into core (#1451)

.changeset/mighty-clocks-share.md

@@ -0,0 +1,6 @@
+---
+"@llamaindex/core": patch
+"llamaindex": patch
+---
+
+refactor: move storage into core

packages/core/package.json

@@ -214,6 +214,48 @@
         "default": "./storage/chat-store/dist/index.js"
       }
     },
+    "./storage/doc-store": {
+      "require": {
+        "types": "./storage/doc-store/dist/index.d.cts",
+        "default": "./storage/doc-store/dist/index.cjs"
+      },
+      "import": {
+        "types": "./storage/doc-store/dist/index.d.ts",
+        "default": "./storage/doc-store/dist/index.js"
+      },
+      "default": {
+        "types": "./storage/doc-store/dist/index.d.ts",
+        "default": "./storage/doc-store/dist/index.js"
+      }
+    },
+    "./storage/index-store": {
+      "require": {
+        "types": "./storage/index-store/dist/index.d.cts",
+        "default": "./storage/index-store/dist/index.cjs"
+      },
+      "import": {
+        "types": "./storage/index-store/dist/index.d.ts",
+        "default": "./storage/index-store/dist/index.js"
+      },
+      "default": {
+        "types": "./storage/index-store/dist/index.d.ts",
+        "default": "./storage/index-store/dist/index.js"
+      }
+    },
+    "./storage/kv-store": {
+      "require": {
+        "types": "./storage/kv-store/dist/index.d.cts",
+        "default": "./storage/kv-store/dist/index.cjs"
+      },
+      "import": {
+        "types": "./storage/kv-store/dist/index.d.ts",
+        "default": "./storage/kv-store/dist/index.js"
+      },
+      "default": {
+        "types": "./storage/kv-store/dist/index.d.ts",
+        "default": "./storage/kv-store/dist/index.js"
+      }
+    },
     "./response-synthesizers": {
       "require": {
         "types": "./response-synthesizers/dist/index.d.cts",

packages/core/src/data-structs/data-structs.ts

@@ -1,5 +1,6 @@
 import { randomUUID } from "@llamaindex/env";
 import type { UUID } from "../global";
+import { BaseNode } from "../schema";
 import { IndexStructType } from "./struct-type";
 
 export abstract class IndexStruct {
@@ -65,3 +66,48 @@ export class KeywordTable extends IndexStruct {
     };
   }
 }
+
+export class IndexDict extends IndexStruct {
+  nodesDict: Record<string, BaseNode> = {};
+  type: IndexStructType = IndexStructType.SIMPLE_DICT;
+
+  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 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,
+    };
+  }
+}

packages/core/src/data-structs/index.ts

@@ -1,2 +1,8 @@
-export { IndexStruct, KeywordTable } from "./data-structs";
+export {
+  IndexDict,
+  IndexList,
+  IndexStruct,
+  KeywordTable,
+} from "./data-structs";
+export { jsonToIndexStruct } from "./json-to-index-struct";
 export { IndexStructType } from "./struct-type";

packages/core/src/data-structs/json-to-index-struct.ts

@@ -0,0 +1,26 @@
+import type { BaseNode } from "../schema";
+import { jsonToNode } from "../schema";
+import { IndexDict, IndexList, IndexStruct } from "./data-structs";
+import { IndexStructType } from "./struct-type";
+
+export function jsonToIndexStruct(
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  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}`);
+  }
+}

packages/core/src/query-engine/retriever.ts

@@ -1,6 +1,5 @@
 import type { MessageContent } from "../llms";
 import type { BaseNodePostprocessor } from "../postprocessor";
-import { BaseQueryEngine, type QueryType } from "../query-engine";
 import {
   type BaseSynthesizer,
   getResponseSynthesizer,
@@ -8,6 +7,7 @@ import {
 import { BaseRetriever } from "../retriever";
 import type { NodeWithScore } from "../schema";
 import { extractText } from "../utils";
+import { BaseQueryEngine, type QueryType } from "./base";
 
 export class RetrieverQueryEngine extends BaseQueryEngine {
   retriever: BaseRetriever;

packages/core/src/schema/index.ts

@@ -1,5 +1,10 @@
 export * from "./node";
-export { FileReader, TransformComponent, type BaseReader } from "./type";
+export {
+  FileReader,
+  TransformComponent,
+  type BaseReader,
+  type StoredValue,
+} from "./type";
 export type { BaseOutputParser } from "./type/base-output-parser";
 export { EngineResponse } from "./type/engine–response";
 export * from "./zod";

packages/core/src/schema/type.ts

@@ -1,6 +1,10 @@
 import { fs, path, randomUUID } from "@llamaindex/env";
 import type { BaseNode, Document } from "./node";
 
+// fixme: remove any
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export type StoredValue = Record<string, any> | null;
+
 interface TransformComponentSignature<
   Result extends BaseNode[] | Promise<BaseNode[]>,
 > {

packages/core/src/storage/doc-store/index.ts

@@ -0,0 +1,167 @@
+import { path } from "@llamaindex/env";
+import {
+  DEFAULT_DOC_STORE_PERSIST_FILENAME,
+  DEFAULT_PERSIST_DIR,
+} from "../../global";
+import type { StoredValue } from "../../schema";
+import { BaseNode, Document, ObjectType, TextNode } from "../../schema";
+
+const TYPE_KEY = "__type__";
+const DATA_KEY = "__data__";
+
+export interface Serializer<T> {
+  toPersistence(data: Record<string, unknown>): T;
+
+  fromPersistence(data: T): Record<string, unknown>;
+}
+
+export const jsonSerializer: Serializer<string> = {
+  toPersistence(data) {
+    return JSON.stringify(data);
+  },
+  fromPersistence(data) {
+    return JSON.parse(data);
+  },
+};
+
+export const noneSerializer: Serializer<Record<string, unknown>> = {
+  toPersistence(data) {
+    return data;
+  },
+  fromPersistence(data) {
+    return data;
+  },
+};
+
+type DocJson<Data> = {
+  [TYPE_KEY]: ObjectType;
+  [DATA_KEY]: Data;
+};
+
+export function isValidDocJson(
+  docJson: StoredValue | null | undefined,
+): docJson is DocJson<unknown> {
+  return (
+    typeof docJson === "object" &&
+    docJson !== null &&
+    docJson[TYPE_KEY] !== undefined &&
+    docJson[DATA_KEY] !== undefined
+  );
+}
+
+export function docToJson(
+  doc: BaseNode,
+  serializer: Serializer<unknown>,
+): DocJson<unknown> {
+  return {
+    [DATA_KEY]: serializer.toPersistence(doc.toJSON()),
+    [TYPE_KEY]: doc.type,
+  };
+}
+
+export function jsonToDoc<Data>(
+  docDict: DocJson<Data>,
+  serializer: Serializer<Data>,
+): BaseNode {
+  const docType = docDict[TYPE_KEY];
+  // fixme: zod type check this
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  const dataDict: any = serializer.fromPersistence(docDict[DATA_KEY]);
+  let doc: BaseNode;
+
+  if (docType === ObjectType.DOCUMENT) {
+    doc = new Document({
+      text: dataDict.text,
+      id_: dataDict.id_,
+      embedding: dataDict.embedding,
+      hash: dataDict.hash,
+      metadata: dataDict.metadata,
+    });
+  } else if (docType === ObjectType.TEXT) {
+    doc = new TextNode({
+      text: dataDict.text,
+      id_: dataDict.id_,
+      hash: dataDict.hash,
+      metadata: dataDict.metadata,
+      relationships: dataDict.relationships,
+    });
+  } else {
+    throw new Error(`Unknown doc type: ${docType}`);
+  }
+
+  return doc;
+}
+
+const DEFAULT_PERSIST_PATH = path.join(
+  DEFAULT_PERSIST_DIR,
+  DEFAULT_DOC_STORE_PERSIST_FILENAME,
+);
+
+export interface RefDocInfo {
+  nodeIds: string[];
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  extraInfo: Record<string, any>;
+}
+
+export abstract class BaseDocumentStore {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  serializer: Serializer<any> = jsonSerializer;
+
+  // Save/load
+  persist(persistPath: string = DEFAULT_PERSIST_PATH): void {
+    // Persist the docstore to a file.
+  }
+
+  // Main interface
+  abstract docs(): Promise<Record<string, BaseNode>>;
+
+  abstract addDocuments(docs: BaseNode[], allowUpdate: boolean): Promise<void>;
+
+  abstract getDocument(
+    docId: string,
+    raiseError: boolean,
+  ): Promise<BaseNode | undefined>;
+
+  abstract deleteDocument(docId: string, raiseError: boolean): Promise<void>;
+
+  abstract documentExists(docId: string): Promise<boolean>;
+
+  // Hash
+  abstract setDocumentHash(docId: string, docHash: string): Promise<void>;
+
+  abstract getDocumentHash(docId: string): Promise<string | undefined>;
+
+  abstract getAllDocumentHashes(): Promise<Record<string, string>>;
+
+  // Ref Docs
+  abstract getAllRefDocInfo(): Promise<Record<string, RefDocInfo> | undefined>;
+
+  abstract getRefDocInfo(refDocId: string): Promise<RefDocInfo | undefined>;
+
+  abstract deleteRefDoc(refDocId: string, raiseError: boolean): Promise<void>;
+
+  // Nodes
+  getNodes(nodeIds: string[], raiseError: boolean = true): Promise<BaseNode[]> {
+    return Promise.all(
+      nodeIds.map((nodeId) => this.getNode(nodeId, raiseError)),
+    );
+  }
+
+  async getNode(nodeId: string, raiseError: boolean = true): Promise<BaseNode> {
+    const doc = await this.getDocument(nodeId, raiseError);
+    if (!(doc instanceof BaseNode)) {
+      throw new Error(`Document ${nodeId} is not a Node.`);
+    }
+    return doc;
+  }
+
+  async getNodeDict(nodeIdDict: {
+    [index: number]: string;
+  }): Promise<Record<number, BaseNode>> {
+    const result: Record<number, BaseNode> = {};
+    for (const index in nodeIdDict) {
+      result[index] = await this.getNode(nodeIdDict[index]!);
+    }
+    return result;
+  }
+}

packages/core/src/storage/index-store/index.ts

@@ -0,0 +1,115 @@
+import { path } from "@llamaindex/env";
+import { IndexStruct, jsonToIndexStruct } from "../../data-structs";
+import {
+  DEFAULT_INDEX_STORE_PERSIST_FILENAME,
+  DEFAULT_NAMESPACE,
+  DEFAULT_PERSIST_DIR,
+} from "../../global";
+import {
+  BaseInMemoryKVStore,
+  BaseKVStore,
+  type DataType,
+  SimpleKVStore,
+} from "../kv-store";
+
+export const DEFAULT_PERSIST_PATH = path.join(
+  DEFAULT_PERSIST_DIR,
+  DEFAULT_INDEX_STORE_PERSIST_FILENAME,
+);
+
+export abstract class BaseIndexStore {
+  abstract getIndexStructs(): Promise<IndexStruct[]>;
+
+  abstract addIndexStruct(indexStruct: IndexStruct): Promise<void>;
+
+  abstract deleteIndexStruct(key: string): Promise<void>;
+
+  abstract getIndexStruct(structId?: string): Promise<IndexStruct | undefined>;
+
+  async persist(persistPath: string = DEFAULT_PERSIST_PATH): Promise<void> {
+    // Persist the index store to disk.
+  }
+}
+
+export class KVIndexStore extends BaseIndexStore {
+  private _kvStore: BaseKVStore;
+  private _collection: string;
+
+  constructor(kvStore: BaseKVStore, namespace: string = DEFAULT_NAMESPACE) {
+    super();
+    this._kvStore = kvStore;
+    this._collection = `${namespace}/data`;
+  }
+
+  async addIndexStruct(indexStruct: IndexStruct): Promise<void> {
+    const key = indexStruct.indexId;
+    const data = indexStruct.toJson();
+    await this._kvStore.put(key, data, this._collection);
+  }
+
+  async deleteIndexStruct(key: string): Promise<void> {
+    await this._kvStore.delete(key, this._collection);
+  }
+
+  async getIndexStruct(structId?: string): Promise<IndexStruct | undefined> {
+    if (!structId) {
+      const structs = await this.getIndexStructs();
+      if (structs.length !== 1) {
+        throw new Error("More than one index struct found");
+      }
+      return structs[0];
+    } else {
+      const json = await this._kvStore.get(structId, this._collection);
+      if (json == null) {
+        return;
+      }
+      return jsonToIndexStruct(json);
+    }
+  }
+
+  async getIndexStructs(): Promise<IndexStruct[]> {
+    const jsons = await this._kvStore.getAll(this._collection);
+    return Object.values(jsons).map((json) => jsonToIndexStruct(json));
+  }
+}
+
+export class SimpleIndexStore extends KVIndexStore {
+  private kvStore: BaseInMemoryKVStore;
+
+  constructor(kvStore?: BaseInMemoryKVStore) {
+    kvStore = kvStore || new SimpleKVStore();
+    super(kvStore);
+    this.kvStore = kvStore;
+  }
+
+  static async fromPersistDir(
+    persistDir: string = DEFAULT_PERSIST_DIR,
+  ): Promise<SimpleIndexStore> {
+    const persistPath = path.join(
+      persistDir,
+      DEFAULT_INDEX_STORE_PERSIST_FILENAME,
+    );
+    return this.fromPersistPath(persistPath);
+  }
+
+  static async fromPersistPath(persistPath: string): Promise<SimpleIndexStore> {
+    const simpleKVStore = await SimpleKVStore.fromPersistPath(persistPath);
+    return new SimpleIndexStore(simpleKVStore);
+  }
+
+  async persist(persistPath: string = DEFAULT_PERSIST_DIR): Promise<void> {
+    this.kvStore.persist(persistPath);
+  }
+
+  static fromDict(saveDict: DataType): SimpleIndexStore {
+    const simpleKVStore = SimpleKVStore.fromDict(saveDict);
+    return new SimpleIndexStore(simpleKVStore);
+  }
+
+  toDict(): Record<string, unknown> {
+    if (!(this.kvStore instanceof SimpleKVStore)) {
+      throw new Error("KVStore is not a SimpleKVStore");
+    }
+    return this.kvStore.toDict();
+  }
+}

packages/core/src/storage/kv-store/index.ts

@@ -1,7 +1,34 @@
-import { DEFAULT_COLLECTION } from "@llamaindex/core/global";
 import { fs, path } from "@llamaindex/env";
-import { exists } from "../FileSystem.js";
-import { BaseKVStore, type StoredValue } from "./types.js";
+
+import { DEFAULT_COLLECTION } from "../../global";
+import type { StoredValue } from "../../schema";
+
+async function exists(path: string): Promise<boolean> {
+  try {
+    await fs.access(path);
+    return true;
+  } catch {
+    return false;
+  }
+}
+
+export abstract class BaseKVStore {
+  abstract put(
+    key: string,
+    val: StoredValue,
+    collection?: string,
+  ): Promise<void>;
+  abstract get(key: string, collection?: string): Promise<StoredValue>;
+  abstract getAll(collection?: string): Promise<Record<string, StoredValue>>;
+  abstract delete(key: string, collection?: string): Promise<boolean>;
+}
+
+export abstract class BaseInMemoryKVStore extends BaseKVStore {
+  abstract persist(persistPath: string): void;
+  static fromPersistPath(persistPath: string): BaseInMemoryKVStore {
+    throw new Error("Method not implemented.");
+  }
+}
 
 export type DataType = Record<string, Record<string, StoredValue>>;
 
@@ -42,8 +69,10 @@ export class SimpleKVStore extends BaseKVStore {
   }
 
   async getAll(collection: string = DEFAULT_COLLECTION) {
-    // fixme: null value here
-    return structuredClone(this.data[collection])!; // Creating a shallow copy of the object
+    if (this.data[collection]) {
+      return structuredClone(this.data[collection]);
+    }
+    return {};
   }
 
   async delete(

packages/core/storage/doc-store/package.json

@@ -0,0 +1,8 @@
+{
+  "type": "module",
+  "main": "./dist/index.cjs",
+  "module": "./dist/index.js",
+  "types": "./dist/index.d.ts",
+  "exports": "./dist/index.js",
+  "private": true
+}

packages/core/storage/index-store/package.json

@@ -0,0 +1,8 @@
+{
+  "type": "module",
+  "main": "./dist/index.cjs",
+  "module": "./dist/index.js",
+  "types": "./dist/index.d.ts",
+  "exports": "./dist/index.js",
+  "private": true
+}

packages/core/storage/kv-store/package.json

@@ -0,0 +1,8 @@
+{
+  "type": "module",
+  "main": "./dist/index.cjs",
+  "module": "./dist/index.js",
+  "types": "./dist/index.d.ts",
+  "exports": "./dist/index.js",
+  "private": true
+}

packages/llamaindex/src/index.edge.ts

@@ -18,6 +18,7 @@ export {
 } from "@llamaindex/cloud/reader";
 export * from "@llamaindex/core/agent";
 export * from "@llamaindex/core/chat-engine";
+export * from "@llamaindex/core/data-structs";
 export {
   CallbackManager,
   DEFAULT_BASE_URL,
@@ -57,6 +58,10 @@ export * from "@llamaindex/core/query-engine";
 export * from "@llamaindex/core/response-synthesizers";
 export * from "@llamaindex/core/retriever";
 export * from "@llamaindex/core/schema";
+export * from "@llamaindex/core/storage/chat-store";
+export * from "@llamaindex/core/storage/doc-store";
+export * from "@llamaindex/core/storage/index-store";
+export * from "@llamaindex/core/storage/kv-store";
 export * from "./agent/index.js";
 export * from "./cloud/index.js";
 export * from "./embeddings/index.js";

packages/llamaindex/src/indices/BaseIndex.ts

@@ -2,12 +2,12 @@ import type { BaseQueryEngine } from "@llamaindex/core/query-engine";
 import type { BaseSynthesizer } from "@llamaindex/core/response-synthesizers";
 import type { BaseRetriever } from "@llamaindex/core/retriever";
 import type { BaseNode, Document } from "@llamaindex/core/schema";
+import type { BaseDocumentStore } from "@llamaindex/core/storage/doc-store";
+import type { BaseIndexStore } from "@llamaindex/core/storage/index-store";
 import type { ServiceContext } from "../ServiceContext.js";
 import { nodeParserFromSettingsOrContext } from "../Settings.js";
 import { runTransformations } from "../ingestion/IngestionPipeline.js";
 import type { StorageContext } from "../storage/StorageContext.js";
-import type { BaseDocumentStore } from "../storage/docStore/types.js";
-import type { BaseIndexStore } from "../storage/indexStore/types.js";
 
 export interface BaseIndexInit<T> {
   serviceContext?: ServiceContext | undefined;

packages/llamaindex/src/indices/IndexStruct.ts

@@ -1,28 +0,0 @@
-import { randomUUID } from "@llamaindex/env";
-
-/**
- * The underlying structure of each index.
- */
-export abstract class IndexStruct {
-  indexId: string;
-  summary?: string | undefined;
-
-  constructor(indexId = randomUUID(), summary: string | undefined = undefined) {
-    this.indexId = indexId;
-    this.summary = summary;
-  }
-
-  toJson(): Record<string, unknown> {
-    return {
-      indexId: this.indexId,
-      summary: this.summary,
-    };
-  }
-
-  getSummary(): string {
-    if (this.summary === undefined) {
-      throw new Error("summary field of the index dict is not set");
-    }
-    return this.summary;
-  }
-}

packages/llamaindex/src/indices/index.ts

@@ -1,6 +1,5 @@
+export * from "@llamaindex/core/indices";
 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";

packages/llamaindex/src/indices/json-to-index-struct.ts

@@ -1,80 +0,0 @@
-import type { BaseNode } from "@llamaindex/core/schema";
-import { jsonToNode } from "@llamaindex/core/schema";
-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];
-  }
-}
-
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-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,
-    };
-  }
-}

packages/llamaindex/src/indices/keyword/index.ts

@@ -10,17 +10,15 @@ import { serviceContextFromDefaults } from "../../ServiceContext.js";
 import { RetrieverQueryEngine } from "../../engines/query/index.js";
 import type { StorageContext } from "../../storage/StorageContext.js";
 import { storageContextFromDefaults } from "../../storage/StorageContext.js";
-import type { BaseDocumentStore } from "../../storage/docStore/types.js";
 import type { BaseIndexInit } from "../BaseIndex.js";
 import { BaseIndex } from "../BaseIndex.js";
-import { IndexStructType } from "../json-to-index-struct.js";
 import {
   extractKeywordsGivenResponse,
   rakeExtractKeywords,
   simpleExtractKeywords,
 } from "./utils.js";
 
-import { KeywordTable } from "@llamaindex/core/data-structs";
+import { IndexStructType, KeywordTable } from "@llamaindex/core/data-structs";
 import type { LLM } from "@llamaindex/core/llms";
 import type { BaseNodePostprocessor } from "@llamaindex/core/postprocessor";
 import {
@@ -34,6 +32,7 @@ import type {
   QueryBundle,
 } from "@llamaindex/core/query-engine";
 import { BaseRetriever } from "@llamaindex/core/retriever";
+import type { BaseDocumentStore } from "@llamaindex/core/storage/doc-store";
 import { extractText } from "@llamaindex/core/utils";
 import { llmFromSettingsOrContext } from "../../Settings.js";
 

packages/llamaindex/src/indices/summary/index.ts

@@ -1,3 +1,4 @@
+import { IndexList, IndexStructType } from "@llamaindex/core/data-structs";
 import type { BaseNodePostprocessor } from "@llamaindex/core/postprocessor";
 import {
   type ChoiceSelectPrompt,
@@ -12,6 +13,10 @@ import type {
   Document,
   NodeWithScore,
 } from "@llamaindex/core/schema";
+import type {
+  BaseDocumentStore,
+  RefDocInfo,
+} from "@llamaindex/core/storage/doc-store";
 import { extractText } from "@llamaindex/core/utils";
 import _ from "lodash";
 import type { ServiceContext } from "../../ServiceContext.js";
@@ -22,13 +27,8 @@ import {
 import { RetrieverQueryEngine } from "../../engines/query/index.js";
 import type { StorageContext } from "../../storage/StorageContext.js";
 import { storageContextFromDefaults } from "../../storage/StorageContext.js";
-import type {
-  BaseDocumentStore,
-  RefDocInfo,
-} from "../../storage/docStore/types.js";
 import type { BaseIndexInit } from "../BaseIndex.js";
 import { BaseIndex } from "../BaseIndex.js";
-import { IndexList, IndexStructType } from "../json-to-index-struct.js";
 import type {
   ChoiceSelectParserFunction,
   NodeFormatterFunction,

packages/llamaindex/src/indices/vectorStore/index.ts

@@ -1,3 +1,4 @@
+import { IndexDict, IndexStructType } from "@llamaindex/core/data-structs";
 import {
   DEFAULT_SIMILARITY_TOP_K,
   type BaseEmbedding,
@@ -16,6 +17,7 @@ import {
   type Document,
   type NodeWithScore,
 } from "@llamaindex/core/schema";
+import type { BaseIndexStore } from "@llamaindex/core/storage/index-store";
 import type { ServiceContext } from "../../ServiceContext.js";
 import { nodeParserFromSettingsOrContext } from "../../Settings.js";
 import { RetrieverQueryEngine } from "../../engines/query/RetrieverQueryEngine.js";
@@ -29,7 +31,6 @@ import {
 } from "../../ingestion/strategies/index.js";
 import type { StorageContext } from "../../storage/StorageContext.js";
 import { storageContextFromDefaults } from "../../storage/StorageContext.js";
-import type { BaseIndexStore } from "../../storage/indexStore/types.js";
 import type {
   BaseVectorStore,
   MetadataFilters,
@@ -39,7 +40,6 @@ import type {
 import { VectorStoreQueryMode } from "../../vector-store/types.js";
 import type { BaseIndexInit } from "../BaseIndex.js";
 import { BaseIndex } from "../BaseIndex.js";
-import { IndexDict, IndexStructType } from "../json-to-index-struct.js";
 
 interface IndexStructOptions {
   indexStruct?: IndexDict | undefined;

packages/llamaindex/src/ingestion/IngestionCache.ts

@@ -1,13 +1,15 @@
 import type { BaseNode, TransformComponent } from "@llamaindex/core/schema";
 import { MetadataMode } from "@llamaindex/core/schema";
-import { createSHA256 } from "@llamaindex/env";
 import {
   docToJson,
   jsonSerializer,
   jsonToDoc,
-} from "../storage/docStore/utils.js";
-import { SimpleKVStore } from "../storage/kvStore/SimpleKVStore.js";
-import type { BaseKVStore } from "../storage/kvStore/types.js";
+} from "@llamaindex/core/storage/doc-store";
+import {
+  SimpleKVStore,
+  type BaseKVStore,
+} from "@llamaindex/core/storage/kv-store";
+import { createSHA256 } from "@llamaindex/env";
 
 const transformToJSON = (obj: TransformComponent) => {
   // eslint-disable-next-line @typescript-eslint/no-explicit-any

packages/llamaindex/src/ingestion/IngestionPipeline.ts

@@ -6,7 +6,7 @@ import {
   type Document,
   type Metadata,
 } from "@llamaindex/core/schema";
-import type { BaseDocumentStore } from "../storage/docStore/types.js";
+import type { BaseDocumentStore } from "@llamaindex/core/storage/doc-store";
 import type {
   BaseVectorStore,
   VectorStoreByType,

packages/llamaindex/src/ingestion/strategies/DuplicatesStrategy.ts

@@ -1,5 +1,5 @@
 import { BaseNode, TransformComponent } from "@llamaindex/core/schema";
-import type { BaseDocumentStore } from "../../storage/docStore/types.js";
+import type { BaseDocumentStore } from "@llamaindex/core/storage/doc-store";
 
 /**
  * Handle doc store duplicates by checking all hashes.

packages/llamaindex/src/ingestion/strategies/UpsertsAndDeleteStrategy.ts

@@ -1,5 +1,5 @@
 import { BaseNode, TransformComponent } from "@llamaindex/core/schema";
-import type { BaseDocumentStore } from "../../storage/docStore/types.js";
+import type { BaseDocumentStore } from "@llamaindex/core/storage/doc-store";
 import type { BaseVectorStore } from "../../vector-store/types.js";
 import { classify } from "./classify.js";
 

packages/llamaindex/src/ingestion/strategies/UpsertsStrategy.ts

@@ -1,5 +1,5 @@
 import { BaseNode, TransformComponent } from "@llamaindex/core/schema";
-import type { BaseDocumentStore } from "../../storage/docStore/types.js";
+import type { BaseDocumentStore } from "@llamaindex/core/storage/doc-store";
 import type { BaseVectorStore } from "../../vector-store/types.js";
 import { classify } from "./classify.js";
 

packages/llamaindex/src/ingestion/strategies/classify.ts

@@ -1,5 +1,5 @@
 import type { BaseNode } from "@llamaindex/core/schema";
-import type { BaseDocumentStore } from "../../storage/docStore/types.js";
+import type { BaseDocumentStore } from "@llamaindex/core/storage/doc-store";
 
 export async function classify(docStore: BaseDocumentStore, nodes: BaseNode[]) {
   const existingDocIds = Object.values(await docStore.getAllDocumentHashes());

packages/llamaindex/src/ingestion/strategies/index.ts

@@ -1,5 +1,5 @@
 import { TransformComponent } from "@llamaindex/core/schema";
-import type { BaseDocumentStore } from "../../storage/docStore/types.js";
+import type { BaseDocumentStore } from "@llamaindex/core/storage/doc-store";
 import type { BaseVectorStore } from "../../vector-store/types.js";
 import { DuplicatesStrategy } from "./DuplicatesStrategy.js";
 import { UpsertsAndDeleteStrategy } from "./UpsertsAndDeleteStrategy.js";

packages/llamaindex/src/storage/StorageContext.ts

@@ -4,6 +4,11 @@ import {
   DEFAULT_NAMESPACE,
 } from "@llamaindex/core/global";
 import { ModalityType, ObjectType } from "@llamaindex/core/schema";
+import type { BaseDocumentStore } from "@llamaindex/core/storage/doc-store";
+import {
+  BaseIndexStore,
+  SimpleIndexStore,
+} from "@llamaindex/core/storage/index-store";
 import { path } from "@llamaindex/env";
 import type { ServiceContext } from "../ServiceContext.js";
 import { SimpleVectorStore } from "../vector-store/SimpleVectorStore.js";
@@ -12,9 +17,6 @@ import type {
   VectorStoreByType,
 } from "../vector-store/types.js";
 import { SimpleDocumentStore } from "./docStore/SimpleDocumentStore.js";
-import type { BaseDocumentStore } from "./docStore/types.js";
-import { SimpleIndexStore } from "./indexStore/SimpleIndexStore.js";
-import type { BaseIndexStore } from "./indexStore/types.js";
 
 export interface StorageContext {
   docStore: BaseDocumentStore;

packages/llamaindex/src/storage/docStore/KVDocumentStore.ts

@@ -1,11 +1,15 @@
 import { DEFAULT_NAMESPACE } from "@llamaindex/core/global";
 import type { BaseNode } from "@llamaindex/core/schema";
 import { ObjectType } from "@llamaindex/core/schema";
+import type { RefDocInfo } from "@llamaindex/core/storage/doc-store";
+import {
+  BaseDocumentStore,
+  docToJson,
+  isValidDocJson,
+  jsonToDoc,
+} from "@llamaindex/core/storage/doc-store";
+import type { BaseKVStore } from "@llamaindex/core/storage/kv-store";
 import _ from "lodash";
-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";
 
 type DocMetaData = { docHash: string; refDocId?: string };
 

packages/llamaindex/src/storage/docStore/PostgresDocumentStore.ts

@@ -1,10 +1,10 @@
 import { DEFAULT_NAMESPACE } from "@llamaindex/core/global";
+import { noneSerializer } from "@llamaindex/core/storage/doc-store";
 import {
   PostgresKVStore,
   type PostgresKVStoreConfig,
 } from "../kvStore/PostgresKVStore.js";
 import { KVDocumentStore } from "./KVDocumentStore.js";
-import { noneSerializer } from "./utils.js";
 
 const DEFAULT_TABLE_NAME = "llamaindex_doc_store";
 

packages/llamaindex/src/storage/docStore/SimpleDocumentStore.ts

@@ -3,10 +3,12 @@ import {
   DEFAULT_NAMESPACE,
   DEFAULT_PERSIST_DIR,
 } from "@llamaindex/core/global";
+import {
+  BaseInMemoryKVStore,
+  SimpleKVStore,
+} from "@llamaindex/core/storage/kv-store";
 import { path } from "@llamaindex/env";
 import _ from "lodash";
-import { SimpleKVStore } from "../kvStore/SimpleKVStore.js";
-import { BaseInMemoryKVStore } from "../kvStore/types.js";
 import { KVDocumentStore } from "./KVDocumentStore.js";
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any

packages/llamaindex/src/storage/docStore/types.ts

@@ -1,77 +0,0 @@
-import {
-  DEFAULT_DOC_STORE_PERSIST_FILENAME,
-  DEFAULT_PERSIST_DIR,
-} from "@llamaindex/core/global";
-import { BaseNode } from "@llamaindex/core/schema";
-import { jsonSerializer, type Serializer } from "./utils.js";
-
-const defaultPersistPath = `${DEFAULT_PERSIST_DIR}/${DEFAULT_DOC_STORE_PERSIST_FILENAME}`;
-
-export interface RefDocInfo {
-  nodeIds: string[];
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  extraInfo: Record<string, any>;
-}
-
-export abstract class BaseDocumentStore {
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  serializer: Serializer<any> = jsonSerializer;
-
-  // Save/load
-  persist(persistPath: string = defaultPersistPath): void {
-    // Persist the docstore to a file.
-  }
-
-  // Main interface
-  abstract docs(): Promise<Record<string, BaseNode>>;
-
-  abstract addDocuments(docs: BaseNode[], allowUpdate: boolean): Promise<void>;
-
-  abstract getDocument(
-    docId: string,
-    raiseError: boolean,
-  ): Promise<BaseNode | undefined>;
-
-  abstract deleteDocument(docId: string, raiseError: boolean): Promise<void>;
-
-  abstract documentExists(docId: string): Promise<boolean>;
-
-  // Hash
-  abstract setDocumentHash(docId: string, docHash: string): Promise<void>;
-
-  abstract getDocumentHash(docId: string): Promise<string | undefined>;
-
-  abstract getAllDocumentHashes(): Promise<Record<string, string>>;
-
-  // Ref Docs
-  abstract getAllRefDocInfo(): Promise<Record<string, RefDocInfo> | undefined>;
-
-  abstract getRefDocInfo(refDocId: string): Promise<RefDocInfo | undefined>;
-
-  abstract deleteRefDoc(refDocId: string, raiseError: boolean): Promise<void>;
-
-  // Nodes
-  getNodes(nodeIds: string[], raiseError: boolean = true): Promise<BaseNode[]> {
-    return Promise.all(
-      nodeIds.map((nodeId) => this.getNode(nodeId, raiseError)),
-    );
-  }
-
-  async getNode(nodeId: string, raiseError: boolean = true): Promise<BaseNode> {
-    const doc = await this.getDocument(nodeId, raiseError);
-    if (!(doc instanceof BaseNode)) {
-      throw new Error(`Document ${nodeId} is not a Node.`);
-    }
-    return doc;
-  }
-
-  async getNodeDict(nodeIdDict: {
-    [index: number]: string;
-  }): Promise<Record<number, BaseNode>> {
-    const result: Record<number, BaseNode> = {};
-    for (const index in nodeIdDict) {
-      result[index] = await this.getNode(nodeIdDict[index]!);
-    }
-    return result;
-  }
-}

packages/llamaindex/src/storage/docStore/utils.ts

@@ -1,88 +0,0 @@
-import type { BaseNode } from "@llamaindex/core/schema";
-import { Document, ObjectType, TextNode } from "@llamaindex/core/schema";
-import type { StoredValue } from "../kvStore/types.js";
-
-const TYPE_KEY = "__type__";
-const DATA_KEY = "__data__";
-
-export interface Serializer<T> {
-  toPersistence(data: Record<string, unknown>): T;
-  fromPersistence(data: T): Record<string, unknown>;
-}
-
-export const jsonSerializer: Serializer<string> = {
-  toPersistence(data) {
-    return JSON.stringify(data);
-  },
-  fromPersistence(data) {
-    return JSON.parse(data);
-  },
-};
-
-export const noneSerializer: Serializer<Record<string, unknown>> = {
-  toPersistence(data) {
-    return data;
-  },
-  fromPersistence(data) {
-    return data;
-  },
-};
-
-type DocJson<Data> = {
-  [TYPE_KEY]: ObjectType;
-  [DATA_KEY]: Data;
-};
-
-export function isValidDocJson(
-  docJson: StoredValue | null | undefined,
-): docJson is DocJson<unknown> {
-  return (
-    typeof docJson === "object" &&
-    docJson !== null &&
-    docJson[TYPE_KEY] !== undefined &&
-    docJson[DATA_KEY] !== undefined
-  );
-}
-
-export function docToJson(
-  doc: BaseNode,
-  serializer: Serializer<unknown>,
-): DocJson<unknown> {
-  return {
-    [DATA_KEY]: serializer.toPersistence(doc.toJSON()),
-    [TYPE_KEY]: doc.type,
-  };
-}
-
-export function jsonToDoc<Data>(
-  docDict: DocJson<Data>,
-  serializer: Serializer<Data>,
-): BaseNode {
-  const docType = docDict[TYPE_KEY];
-  // fixme: zod type check this
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  const dataDict: any = serializer.fromPersistence(docDict[DATA_KEY]);
-  let doc: BaseNode;
-
-  if (docType === ObjectType.DOCUMENT) {
-    doc = new Document({
-      text: dataDict.text,
-      id_: dataDict.id_,
-      embedding: dataDict.embedding,
-      hash: dataDict.hash,
-      metadata: dataDict.metadata,
-    });
-  } else if (docType === ObjectType.TEXT) {
-    doc = new TextNode({
-      text: dataDict.text,
-      id_: dataDict.id_,
-      hash: dataDict.hash,
-      metadata: dataDict.metadata,
-      relationships: dataDict.relationships,
-    });
-  } else {
-    throw new Error(`Unknown doc type: ${docType}`);
-  }
-
-  return doc;
-}

packages/llamaindex/src/storage/index.ts

@@ -1,15 +1,13 @@
 export * from "@llamaindex/core/storage/chat-store";
+export * from "@llamaindex/core/storage/doc-store";
+export * from "@llamaindex/core/storage/index-store";
+export * from "@llamaindex/core/storage/kv-store";
 export * from "./docStore/AzureCosmosNoSqlDocumentStore.js";
 export { PostgresDocumentStore } from "./docStore/PostgresDocumentStore.js";
 export { SimpleDocumentStore } from "./docStore/SimpleDocumentStore.js";
-export * from "./docStore/types.js";
 export * from "./FileSystem.js";
 export * from "./indexStore/AzureCosmosNoSqlIndexStore.js";
 export { PostgresIndexStore } from "./indexStore/PostgresIndexStore.js";
-export { SimpleIndexStore } from "./indexStore/SimpleIndexStore.js";
-export * from "./indexStore/types.js";
 export * from "./kvStore/AzureCosmosNoSqlKVStore.js";
 export { PostgresKVStore } from "./kvStore/PostgresKVStore.js";
-export { SimpleKVStore } from "./kvStore/SimpleKVStore.js";
-export * from "./kvStore/types.js";
 export * from "./StorageContext.js";

packages/llamaindex/src/storage/indexStore/KVIndexStore.ts

@@ -1,9 +1,11 @@
+import {
+  type IndexStruct,
+  jsonToIndexStruct,
+} from "@llamaindex/core/data-structs";
 import { DEFAULT_NAMESPACE } from "@llamaindex/core/global";
+import { BaseIndexStore } from "@llamaindex/core/storage/index-store";
+import type { BaseKVStore } from "@llamaindex/core/storage/kv-store";
 import _ from "lodash";
-import type { IndexStruct } from "../../indices/IndexStruct.js";
-import { jsonToIndexStruct } from "../../indices/json-to-index-struct.js";
-import type { BaseKVStore } from "../kvStore/types.js";
-import { BaseIndexStore } from "./types.js";
 
 export class KVIndexStore extends BaseIndexStore {
   private _kvStore: BaseKVStore;

packages/llamaindex/src/storage/indexStore/SimpleIndexStore.ts

@@ -1,50 +0,0 @@
-import {
-  DEFAULT_INDEX_STORE_PERSIST_FILENAME,
-  DEFAULT_PERSIST_DIR,
-} from "@llamaindex/core/global";
-import { path } from "@llamaindex/env";
-import type { DataType } from "../kvStore/SimpleKVStore.js";
-import { SimpleKVStore } from "../kvStore/SimpleKVStore.js";
-import type { BaseInMemoryKVStore } from "../kvStore/types.js";
-import { KVIndexStore } from "./KVIndexStore.js";
-
-export class SimpleIndexStore extends KVIndexStore {
-  private kvStore: BaseInMemoryKVStore;
-
-  constructor(kvStore?: BaseInMemoryKVStore) {
-    kvStore = kvStore || new SimpleKVStore();
-    super(kvStore);
-    this.kvStore = kvStore;
-  }
-
-  static async fromPersistDir(
-    persistDir: string = DEFAULT_PERSIST_DIR,
-  ): Promise<SimpleIndexStore> {
-    const persistPath = path.join(
-      persistDir,
-      DEFAULT_INDEX_STORE_PERSIST_FILENAME,
-    );
-    return this.fromPersistPath(persistPath);
-  }
-
-  static async fromPersistPath(persistPath: string): Promise<SimpleIndexStore> {
-    const simpleKVStore = await SimpleKVStore.fromPersistPath(persistPath);
-    return new SimpleIndexStore(simpleKVStore);
-  }
-
-  async persist(persistPath: string = DEFAULT_PERSIST_DIR): Promise<void> {
-    this.kvStore.persist(persistPath);
-  }
-
-  static fromDict(saveDict: DataType): SimpleIndexStore {
-    const simpleKVStore = SimpleKVStore.fromDict(saveDict);
-    return new SimpleIndexStore(simpleKVStore);
-  }
-
-  toDict(): Record<string, unknown> {
-    if (!(this.kvStore instanceof SimpleKVStore)) {
-      throw new Error("KVStore is not a SimpleKVStore");
-    }
-    return this.kvStore.toDict();
-  }
-}

packages/llamaindex/src/storage/indexStore/types.ts

@@ -1,25 +0,0 @@
-import {
-  DEFAULT_INDEX_STORE_PERSIST_FILENAME,
-  DEFAULT_PERSIST_DIR,
-} from "@llamaindex/core/global";
-import { path } from "@llamaindex/env";
-import type { IndexStruct } from "../../indices/IndexStruct.js";
-
-const defaultPersistPath = path.join(
-  DEFAULT_PERSIST_DIR,
-  DEFAULT_INDEX_STORE_PERSIST_FILENAME,
-);
-
-export abstract class BaseIndexStore {
-  abstract getIndexStructs(): Promise<IndexStruct[]>;
-
-  abstract addIndexStruct(indexStruct: IndexStruct): Promise<void>;
-
-  abstract deleteIndexStruct(key: string): Promise<void>;
-
-  abstract getIndexStruct(structId?: string): Promise<IndexStruct | undefined>;
-
-  async persist(persistPath: string = defaultPersistPath): Promise<void> {
-    // Persist the index store to disk.
-  }
-}

packages/llamaindex/src/storage/kvStore/AzureCosmosNoSqlKVStore.ts

@@ -1,8 +1,8 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 import { Container, CosmosClient, Database } from "@azure/cosmos";
 import { DefaultAzureCredential, type TokenCredential } from "@azure/identity";
+import { BaseKVStore } from "@llamaindex/core/storage/kv-store";
 import { getEnv } from "@llamaindex/env";
-import { BaseKVStore } from "./types.js";
 const USER_AGENT_SUFFIX = "LlamaIndex-CDBNoSQL-KVStore-JavaScript";
 const DEFAULT_CHAT_DATABASE = "KVStoreDB";
 const DEFAULT_CHAT_CONTAINER = "KVStoreContainer";

packages/llamaindex/src/storage/kvStore/PostgresKVStore.ts

@@ -1,6 +1,7 @@
 import { DEFAULT_COLLECTION } from "@llamaindex/core/global";
+import type { StoredValue } from "@llamaindex/core/schema";
+import { BaseKVStore } from "@llamaindex/core/storage/kv-store";
 import type pg from "pg";
-import { BaseKVStore, type StoredValue } from "./types.js";
 
 export type DataType = Record<string, Record<string, StoredValue>>;
 

packages/llamaindex/src/storage/kvStore/types.ts

@@ -1,23 +0,0 @@
-const defaultCollection = "data";
-
-// fixme: remove any
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-export type StoredValue = Record<string, any> | null;
-
-export abstract class BaseKVStore {
-  abstract put(
-    key: string,
-    val: StoredValue,
-    collection?: string,
-  ): Promise<void>;
-  abstract get(key: string, collection?: string): Promise<StoredValue>;
-  abstract getAll(collection?: string): Promise<Record<string, StoredValue>>;
-  abstract delete(key: string, collection?: string): Promise<boolean>;
-}
-
-export abstract class BaseInMemoryKVStore extends BaseKVStore {
-  abstract persist(persistPath: string): void;
-  static fromPersistPath(persistPath: string): BaseInMemoryKVStore {
-    throw new Error("Method not implemented.");
-  }
-}