@mastra/rag 0.1.20-alpha.0 → 0.1.20-alpha.2

package/src/document/schema/node.ts

@@ -0,0 +1,187 @@
+import { createHash, randomUUID } from 'crypto';
+import { NodeRelationship, ObjectType } from './types';
+import type { Metadata, RelatedNodeInfo, RelatedNodeType, BaseNodeParams, TextNodeParams } from './types';
+
+/**
+ * Generic abstract class for retrievable nodes
+ */
+export abstract class BaseNode<T extends Metadata = Metadata> {
+  id_: string;
+  metadata: T;
+  relationships: Partial<Record<NodeRelationship, RelatedNodeType<T>>>;
+
+  @lazyInitHash
+  accessor hash: string = '';
+
+  protected constructor(init?: BaseNodeParams<T>) {
+    const { id_, metadata, relationships } = init || {};
+    this.id_ = id_ ?? randomUUID();
+    this.metadata = metadata ?? ({} as T);
+    this.relationships = relationships ?? {};
+  }
+
+  abstract get type(): ObjectType;
+
+  abstract getContent(): string;
+
+  abstract getMetadataStr(): string;
+
+  get sourceNode(): RelatedNodeInfo<T> | undefined {
+    const relationship = this.relationships[NodeRelationship.SOURCE];
+
+    if (Array.isArray(relationship)) {
+      throw new Error('Source object must be a single RelatedNodeInfo object');
+    }
+
+    return relationship;
+  }
+
+  get prevNode(): RelatedNodeInfo<T> | undefined {
+    const relationship = this.relationships[NodeRelationship.PREVIOUS];
+
+    if (Array.isArray(relationship)) {
+      throw new Error('Previous object must be a single RelatedNodeInfo object');
+    }
+
+    return relationship;
+  }
+
+  get nextNode(): RelatedNodeInfo<T> | undefined {
+    const relationship = this.relationships[NodeRelationship.NEXT];
+
+    if (Array.isArray(relationship)) {
+      throw new Error('Next object must be a single RelatedNodeInfo object');
+    }
+
+    return relationship;
+  }
+
+  get parentNode(): RelatedNodeInfo<T> | undefined {
+    const relationship = this.relationships[NodeRelationship.PARENT];
+
+    if (Array.isArray(relationship)) {
+      throw new Error('Parent object must be a single RelatedNodeInfo object');
+    }
+
+    return relationship;
+  }
+
+  get childNodes(): RelatedNodeInfo<T>[] | undefined {
+    const relationship = this.relationships[NodeRelationship.CHILD];
+
+    if (!Array.isArray(relationship)) {
+      throw new Error('Child object must be a an array of RelatedNodeInfo objects');
+    }
+
+    return relationship;
+  }
+
+  abstract generateHash(): string;
+}
+
+/**
+ * TextNode is the default node type for text.
+ */
+export class TextNode<T extends Metadata = Metadata> extends BaseNode<T> {
+  text: string;
+
+  startCharIdx?: number;
+  endCharIdx?: number;
+  metadataSeparator: string;
+
+  constructor(init: TextNodeParams<T> = {}) {
+    super(init);
+    const { text, startCharIdx, endCharIdx, metadataSeparator } = init;
+    this.text = text ?? '';
+    if (startCharIdx) {
+      this.startCharIdx = startCharIdx;
+    }
+    if (endCharIdx) {
+      this.endCharIdx = endCharIdx;
+    }
+    this.metadataSeparator = metadataSeparator ?? '\n';
+  }
+
+  /**
+   * Generate a hash of the text node.
+   * The ID is not part of the hash as it can change independent of content.
+   * @returns
+   */
+  generateHash() {
+    const hashFunction = createSHA256();
+    hashFunction.update(`type=${this.type}`);
+    hashFunction.update(`startCharIdx=${this.startCharIdx} endCharIdx=${this.endCharIdx}`);
+    hashFunction.update(this.getContent());
+    return hashFunction.digest();
+  }
+
+  get type() {
+    return ObjectType.TEXT;
+  }
+
+  getContent(): string {
+    const metadataStr = this.getMetadataStr().trim();
+    return `${metadataStr}\n\n${this.text}`.trim();
+  }
+
+  getMetadataStr(): string {
+    const usableMetadataKeys = new Set(Object.keys(this.metadata).sort());
+
+    return [...usableMetadataKeys].map(key => `${key}: ${this.metadata[key]}`).join(this.metadataSeparator);
+  }
+
+  getNodeInfo() {
+    return { start: this.startCharIdx, end: this.endCharIdx };
+  }
+
+  getText() {
+    return this.text;
+  }
+}
+
+/**
+ * A document is just a special text node with a docId.
+ */
+export class Document<T extends Metadata = Metadata> extends TextNode<T> {
+  constructor(init?: TextNodeParams<T>) {
+    super(init);
+  }
+
+  get type() {
+    return ObjectType.DOCUMENT;
+  }
+}
+
+function lazyInitHash(
+  value: ClassAccessorDecoratorTarget<BaseNode, string>,
+  _context: ClassAccessorDecoratorContext,
+): ClassAccessorDecoratorResult<BaseNode, string> {
+  return {
+    get() {
+      const oldValue = value.get.call(this);
+      if (oldValue === '') {
+        const hash = this.generateHash();
+        value.set.call(this, hash);
+      }
+      return value.get.call(this);
+    },
+    set(newValue: string) {
+      value.set.call(this, newValue);
+    },
+    init(value: string): string {
+      return value;
+    },
+  };
+}
+
+function createSHA256() {
+  const hash = createHash('sha256');
+  return {
+    update(data: string | Uint8Array): void {
+      hash.update(data);
+    },
+    digest() {
+      return hash.digest('base64');
+    },
+  };
+}

package/src/document/schema/types.ts

@@ -0,0 +1,40 @@
+export enum NodeRelationship {
+  SOURCE = 'SOURCE',
+  PREVIOUS = 'PREVIOUS',
+  NEXT = 'NEXT',
+  PARENT = 'PARENT',
+  CHILD = 'CHILD',
+}
+
+export enum ObjectType {
+  TEXT = 'TEXT',
+  IMAGE = 'IMAGE',
+  INDEX = 'INDEX',
+  DOCUMENT = 'DOCUMENT',
+  IMAGE_DOCUMENT = 'IMAGE_DOCUMENT',
+}
+
+export type Metadata = Record<string, any>;
+
+export interface RelatedNodeInfo<T extends Metadata = Metadata> {
+  nodeId: string;
+  nodeType?: ObjectType;
+  metadata: T;
+  hash?: string;
+}
+
+export type RelatedNodeType<T extends Metadata = Metadata> = RelatedNodeInfo<T> | RelatedNodeInfo<T>[];
+
+export type BaseNodeParams<T extends Metadata = Metadata> = {
+  id_?: string | undefined;
+  metadata?: T | undefined;
+  relationships?: Partial<Record<NodeRelationship, RelatedNodeType<T>>> | undefined;
+  hash?: string | undefined;
+};
+
+export type TextNodeParams<T extends Metadata = Metadata> = BaseNodeParams<T> & {
+  text?: string | undefined;
+  startCharIdx?: number | undefined;
+  endCharIdx?: number | undefined;
+  metadataSeparator?: string | undefined;
+};

package/src/document/transformers/html.ts

@@ -1,5 +1,5 @@
-import { Document } from 'llamaindex';
 import { parse } from 'node-html-better-parser';
+import { Document } from '../schema';
 
 import { RecursiveCharacterTransformer } from './character';
 

package/src/document/transformers/json.ts

@@ -1,4 +1,4 @@
-import { Document } from 'llamaindex';
+import { Document } from '../schema';
 
 export class RecursiveJsonTransformer {
   private maxSize: number;

package/src/document/transformers/markdown.ts

@@ -1,4 +1,4 @@
-import { Document } from 'llamaindex';
+import { Document } from '../schema';
 
 import { Language } from '../types';
 

package/src/document/transformers/text.ts

@@ -1,4 +1,4 @@
-import { Document } from 'llamaindex';
+import { Document } from '../schema';
 
 import type { ChunkOptions } from '../types';
 

package/src/document/transformers/transformer.ts

@@ -1,4 +1,4 @@
-import type { Document } from 'llamaindex';
+import type { Document } from '../schema';
 
 export interface Transformer {
   transformDocuments(documents: Document[]): Document[];