diff --git a/packages/core/src/Document.ts b/packages/core/src/Document.ts index 0c1618cbe2c11900d9a2e8529d80b225afe0cbbf..2b2fc3a872a740a8fabec06402e2d26360f081e0 100644 --- a/packages/core/src/Document.ts +++ b/packages/core/src/Document.ts @@ -1,4 +1,12 @@ import { v4 as uuidv4 } from "uuid"; + +export enum NodeType { + DOCUMENT, + TEXT, + IMAGE, + INDEX, +} + export abstract class BaseDocument { text: string; docId?: string; @@ -45,14 +53,20 @@ export abstract class BaseDocument { getDocHash() { return this.docHash; } + + abstract getType(): NodeType; } export class Document extends BaseDocument { - static getType() { - return "Document"; + getType() { + return NodeType.DOCUMENT; } } export class ImageDocument extends Document { image?: string; + + getType() { + return NodeType.IMAGE; + } } \ No newline at end of file diff --git a/packages/core/src/Node.ts b/packages/core/src/Node.ts index bc1872e807673ae4fc1f5b66668db42338a054ea..8247e104b7d63f8581c6fc92a8920368a8c03587 100644 --- a/packages/core/src/Node.ts +++ b/packages/core/src/Node.ts @@ -1,4 +1,4 @@ -import { BaseDocument } from "./Document"; +import { BaseDocument, NodeType } from "./Document"; export enum DocumentRelationship { SOURCE = "source", @@ -8,12 +8,6 @@ export enum DocumentRelationship { CHILD = "child", } -export enum NodeType { - TEXT, - IMAGE, - INDEX, -} - export class Node extends BaseDocument { relationships: { [key in DocumentRelationship]: string | string[] | null }; @@ -61,6 +55,10 @@ export class Node extends BaseDocument { childNodeIds(): string[] { return []; } + + getType() { + return NodeType.TEXT; + } } export interface NodeWithEmbedding { diff --git a/packages/core/src/storage/constants.ts b/packages/core/src/storage/constants.ts index e00fa7bc128296c53ce2b9e22b7129f3234fdc05..a0ded7c10f1218b27e5f9ebb88e069c35a83dc8c 100644 --- a/packages/core/src/storage/constants.ts +++ b/packages/core/src/storage/constants.ts @@ -5,6 +5,4 @@ export const DEFAULT_DOC_STORE_PERSIST_FILENAME = "docstore.json"; export const DEFAULT_VECTOR_STORE_PERSIST_FILENAME = "vector_store.json"; export const DEFAULT_GRAPH_STORE_PERSIST_FILENAME = "graph_store.json"; export const DEFAULT_NAMESPACE = "docstore"; -export const TYPE_KEY = "__type__"; -export const DATA_KEY = "__data__"; -export { DEFAULT_FS } from "./FileSystem"; \ No newline at end of file +export { DEFAULT_FS } from "./FileSystem"; diff --git a/packages/core/src/storage/docStore/utils.ts b/packages/core/src/storage/docStore/utils.ts index e7667ed5b2d1867fd0aacdc1a9224bd11ec51bd9..059174429a12cb462fd898a5bce86cd5632bdaf9 100644 --- a/packages/core/src/storage/docStore/utils.ts +++ b/packages/core/src/storage/docStore/utils.ts @@ -1,6 +1,8 @@ import { Node } from "../../Node"; -import { BaseDocument, NodeType, Document } from '../../Document'; -import { DATA_KEY, TYPE_KEY } from '../constants'; +import { BaseDocument, Document, NodeType } from '../../Document'; + +const TYPE_KEY = "__type__"; +const DATA_KEY = "__data__"; export function docToJson(doc: BaseDocument): Record<string, any> { @@ -16,9 +18,11 @@ export function jsonToDoc(docDict: Record<string, any>): BaseDocument { let doc: BaseDocument; if (docType === NodeType.DOCUMENT) { - doc = new Document(dataDict.docId, dataDict.text); + doc = new Document(dataDict.text, dataDict.docId, dataDict.embedding, dataDict.docHash); } else if (docType === NodeType.TEXT) { - doc = new Node(dataDict.relationships); + const reslationships = dataDict.relationships; + doc = new Node(reslationships.text, reslationships.docId, + reslationships.embedding, reslationships.docHash); } else { throw new Error(`Unknown doc type: ${docType}`); }