@mastra/rag 1.2.2 → 1.2.3-alpha.1

package/src/document/types.ts

@@ -1,145 +0,0 @@
-import type { TiktokenEncoding, TiktokenModel } from 'js-tiktoken';
-import type {
-  TitleExtractorsArgs,
-  SummaryExtractArgs,
-  QuestionAnswerExtractArgs,
-  KeywordExtractArgs,
-} from './extractors';
-
-export enum Language {
-  CPP = 'cpp',
-  GO = 'go',
-  JAVA = 'java',
-  KOTLIN = 'kotlin',
-  JS = 'js',
-  TS = 'ts',
-  PHP = 'php',
-  PROTO = 'proto',
-  PYTHON = 'python',
-  RST = 'rst',
-  RUBY = 'ruby',
-  RUST = 'rust',
-  SCALA = 'scala',
-  SWIFT = 'swift',
-  MARKDOWN = 'markdown',
-  LATEX = 'latex',
-  HTML = 'html',
-  SOL = 'sol',
-  CSHARP = 'csharp',
-  COBOL = 'cobol',
-  C = 'c',
-  LUA = 'lua',
-  PERL = 'perl',
-  HASKELL = 'haskell',
-  ELIXIR = 'elixir',
-  POWERSHELL = 'powershell',
-}
-
-export type ExtractParams = {
-  title?: TitleExtractorsArgs | boolean;
-  summary?: SummaryExtractArgs | boolean;
-  questions?: QuestionAnswerExtractArgs | boolean;
-  keywords?: KeywordExtractArgs | boolean;
-};
-
-export type BaseChunkOptions = {
-  /**
-   * @deprecated Use `maxSize` instead. Will be removed in next major version.
-   */
-  size?: number;
-  maxSize?: number;
-  overlap?: number;
-  lengthFunction?: (text: string) => number;
-  keepSeparator?: boolean | 'start' | 'end';
-  addStartIndex?: boolean;
-  stripWhitespace?: boolean;
-};
-
-export type CharacterChunkOptions = BaseChunkOptions & {
-  separator?: string;
-  isSeparatorRegex?: boolean;
-};
-
-export type RecursiveChunkOptions = BaseChunkOptions & {
-  separators?: string[];
-  isSeparatorRegex?: boolean;
-  language?: Language;
-};
-
-export type TokenChunkOptions = BaseChunkOptions & {
-  encodingName?: TiktokenEncoding;
-  modelName?: TiktokenModel;
-  allowedSpecial?: Set<string> | 'all';
-  disallowedSpecial?: Set<string> | 'all';
-};
-
-export type MarkdownChunkOptions = BaseChunkOptions & {
-  headers?: [string, string][];
-  returnEachLine?: boolean;
-  stripHeaders?: boolean;
-};
-
-export type SemanticMarkdownChunkOptions = BaseChunkOptions & {
-  joinThreshold?: number;
-  encodingName?: TiktokenEncoding;
-  modelName?: TiktokenModel;
-  allowedSpecial?: Set<string> | 'all';
-  disallowedSpecial?: Set<string> | 'all';
-};
-
-export type HTMLChunkOptions = BaseChunkOptions &
-  (
-    | { headers: [string, string][]; sections?: never; returnEachLine?: boolean }
-    | { sections: [string, string][]; headers?: never }
-  ) & { returnEachLine?: boolean };
-
-export type JsonChunkOptions = BaseChunkOptions & {
-  minSize?: number;
-  ensureAscii?: boolean;
-  convertLists?: boolean;
-};
-
-export type LatexChunkOptions = BaseChunkOptions & {};
-
-export type SentenceChunkOptions = BaseChunkOptions & {
-  maxSize: number; // Override to make required for sentence strategy
-  minSize?: number;
-  targetSize?: number;
-  sentenceEnders?: string[];
-  fallbackToWords?: boolean;
-  fallbackToCharacters?: boolean;
-};
-
-export type StrategyOptions = {
-  recursive: RecursiveChunkOptions;
-  character: CharacterChunkOptions;
-  token: TokenChunkOptions;
-  markdown: MarkdownChunkOptions;
-  html: HTMLChunkOptions;
-  json: JsonChunkOptions;
-  latex: LatexChunkOptions;
-  sentence: SentenceChunkOptions;
-  'semantic-markdown': SemanticMarkdownChunkOptions;
-};
-
-export type ChunkStrategy =
-  | 'recursive'
-  | 'character'
-  | 'token'
-  | 'markdown'
-  | 'html'
-  | 'json'
-  | 'latex'
-  | 'sentence'
-  | 'semantic-markdown';
-
-export type ChunkParams =
-  | ({ strategy?: 'character' } & CharacterChunkOptions & { extract?: ExtractParams })
-  | ({ strategy: 'recursive' } & RecursiveChunkOptions & { extract?: ExtractParams })
-  | ({ strategy: 'token' } & TokenChunkOptions & { extract?: ExtractParams })
-  | ({ strategy: 'markdown' } & MarkdownChunkOptions & { extract?: ExtractParams })
-  | ({ strategy: 'html' } & HTMLChunkOptions & { extract?: ExtractParams })
-  | ({ strategy: 'json' } & JsonChunkOptions & { extract?: ExtractParams })
-  | ({ strategy: 'latex' } & LatexChunkOptions & { extract?: ExtractParams })
-  | ({ strategy: 'sentence' } & SentenceChunkOptions & { extract?: ExtractParams })
-  | ({ strategy: 'semantic-markdown' } & SemanticMarkdownChunkOptions & { extract?: ExtractParams });

package/src/document/validation.ts

@@ -1,158 +0,0 @@
-import { z } from 'zod';
-import type { ChunkStrategy } from './types';
-
-function handleDeprecatedSize<T extends { size?: number; maxSize?: number }>(data: T): Omit<T, 'size'> {
-  if (data.size !== undefined) {
-    console.warn(
-      '[DEPRECATION] `size` is deprecated. Use `maxSize` instead. This will be removed in the next major version.',
-    );
-
-    if (data.maxSize === undefined) {
-      data.maxSize = data.size;
-    }
-  }
-
-  const { size, ...rest } = data;
-  return rest;
-}
-
-// Base options that apply to all strategies
-const baseChunkOptionsSchema = z.object({
-  size: z.number().positive().optional(),
-  maxSize: z.number().positive().optional(),
-  overlap: z.number().min(0).optional(),
-  lengthFunction: z.function().optional(),
-  keepSeparator: z.union([z.boolean(), z.literal('start'), z.literal('end')]).optional(),
-  addStartIndex: z.boolean().optional(),
-  stripWhitespace: z.boolean().optional(),
-});
-
-// Strategy-specific schemas
-const characterChunkOptionsSchema = baseChunkOptionsSchema
-  .extend({
-    separator: z.string().optional(),
-    isSeparatorRegex: z.boolean().optional(),
-  })
-  .strict();
-
-const recursiveChunkOptionsSchema = baseChunkOptionsSchema
-  .extend({
-    separators: z.array(z.string()).optional(),
-    isSeparatorRegex: z.boolean().optional(),
-    language: z.string().optional(),
-  })
-  .strict();
-
-const sentenceChunkOptionsSchema = baseChunkOptionsSchema
-  .extend({
-    maxSize: z.number().positive(),
-    minSize: z.number().positive().optional(),
-    targetSize: z.number().positive().optional(),
-    sentenceEnders: z.array(z.string()).optional(),
-    fallbackToWords: z.boolean().optional(),
-    fallbackToCharacters: z.boolean().optional(),
-  })
-  .strict();
-
-// Predicate to check for Set-like objects
-const isSetLike = (value: unknown): value is Set<any> => {
-  return (
-    typeof value === 'object' &&
-    value !== null &&
-    typeof (value as Set<any>).has === 'function' &&
-    typeof (value as Set<any>).add === 'function' &&
-    typeof (value as Set<any>).delete === 'function' &&
-    typeof (value as Set<any>).clear === 'function' &&
-    typeof (value as Set<any>).size === 'number'
-  );
-};
-
-// Zod schema for a Set or the literal 'all'
-const setOrAllSchema = z
-  .any()
-  .refine(value => value === 'all' || isSetLike(value), {
-    message: "Must be a Set object or the literal 'all'",
-  })
-  .optional();
-
-const tokenChunkOptionsSchema = baseChunkOptionsSchema
-  .extend({
-    encodingName: z.string().optional(),
-    modelName: z.string().optional(),
-    allowedSpecial: setOrAllSchema,
-    disallowedSpecial: setOrAllSchema,
-  })
-  .strict();
-
-const jsonChunkOptionsSchema = baseChunkOptionsSchema
-  .extend({
-    minSize: z.number().positive().optional(),
-    ensureAscii: z.boolean().optional(),
-    convertLists: z.boolean().optional(),
-  })
-  .strict();
-
-const htmlChunkOptionsSchema = baseChunkOptionsSchema
-  .extend({
-    headers: z.array(z.tuple([z.string(), z.string()])).optional(),
-    sections: z.array(z.tuple([z.string(), z.string()])).optional(),
-    returnEachLine: z.boolean().optional(),
-  })
-  .strict();
-
-const markdownChunkOptionsSchema = baseChunkOptionsSchema
-  .extend({
-    headers: z.array(z.tuple([z.string(), z.string()])).optional(),
-    returnEachLine: z.boolean().optional(),
-    stripHeaders: z.boolean().optional(),
-  })
-  .strict();
-
-const semanticMarkdownChunkOptionsSchema = baseChunkOptionsSchema
-  .extend({
-    joinThreshold: z.number().positive().optional(),
-    encodingName: z.string().optional(),
-    modelName: z.string().optional(),
-    allowedSpecial: setOrAllSchema,
-    disallowedSpecial: setOrAllSchema,
-  })
-  .strict();
-
-const latexChunkOptionsSchema = baseChunkOptionsSchema.strict();
-
-// Strategy-specific validation schemas
-const validationSchemas = {
-  character: characterChunkOptionsSchema.transform(handleDeprecatedSize),
-  recursive: recursiveChunkOptionsSchema.transform(handleDeprecatedSize),
-  sentence: sentenceChunkOptionsSchema.transform(handleDeprecatedSize),
-  token: tokenChunkOptionsSchema.transform(handleDeprecatedSize),
-  json: jsonChunkOptionsSchema.transform(handleDeprecatedSize),
-  html: htmlChunkOptionsSchema.transform(handleDeprecatedSize),
-  markdown: markdownChunkOptionsSchema.transform(handleDeprecatedSize),
-  'semantic-markdown': semanticMarkdownChunkOptionsSchema.transform(handleDeprecatedSize),
-  latex: latexChunkOptionsSchema.transform(handleDeprecatedSize),
-} as const;
-
-export function validateChunkParams(strategy: ChunkStrategy, params: any): void {
-  const schema = validationSchemas[strategy];
-  if (!schema) {
-    throw new Error(`Unknown chunking strategy: ${strategy}`);
-  }
-
-  const result = schema.safeParse(params);
-  if (!result.success) {
-    // Extract unrecognized keys for cleaner error message
-    const unrecognizedError = result.error.errors.find((e: any) => e.code === 'unrecognized_keys');
-    if (unrecognizedError && 'keys' in unrecognizedError) {
-      const keys = (unrecognizedError as any).keys.join(', ');
-      throw new Error(`Invalid parameters for ${strategy} strategy: '${keys}' not supported`);
-    }
-
-    // Fallback to general error message for other validation issues
-    const errorMessage = result.error.errors
-      .map((e: any) => `${e.path.length > 0 ? e.path.join('.') : 'parameter'}: ${e.message}`)
-      .join(', ');
-
-    throw new Error(`Invalid parameters for ${strategy} strategy: ${errorMessage}`);
-  }
-}

package/src/graph-rag/index.test.ts

@@ -1,235 +0,0 @@
-import { describe, it, expect, vi, beforeEach } from 'vitest';
-
-import type { GraphChunk, GraphEdge, GraphEmbedding, GraphNode } from './';
-import { GraphRAG } from './';
-
-describe('GraphRAG', () => {
-  beforeEach(() => {
-    vi.clearAllMocks(); // Clear any mock state before each test
-  });
-
-  describe('addNode', () => {
-    it('should throw an error if node does not have an embedding', () => {
-      const graph = new GraphRAG();
-      const node = {
-        id: '1',
-        content: 'Node 1',
-      };
-      expect(() => graph.addNode(node)).toThrow('Node must have an embedding');
-    });
-
-    it('should throw an error if node embedding dimension is not equal to the graph dimension', () => {
-      const graph = new GraphRAG(2);
-      const node: GraphNode = {
-        id: '1',
-        content: 'Node 1',
-        embedding: [1, 2, 3],
-      };
-      expect(() => graph.addNode(node)).toThrow('Embedding dimension must be 2');
-    });
-
-    it('should add a node to the graph', () => {
-      const graph = new GraphRAG(3);
-      const node = {
-        id: '1',
-        content: 'Node 1',
-        embedding: [1, 2, 3],
-      };
-      graph.addNode(node);
-      expect(graph['nodes'].size).toBe(1);
-    });
-  });
-
-  describe('addEdge', () => {
-    it('should throw an error if either source or target node does not exist', () => {
-      const graph = new GraphRAG();
-      const edge: GraphEdge = {
-        source: '1',
-        target: '2',
-        weight: 0.5,
-        type: 'semantic',
-      };
-      expect(() => graph.addEdge(edge)).toThrow('Both source and target nodes must exist');
-    });
-
-    it('should add an edge between two nodes', () => {
-      const graph = new GraphRAG(3);
-      const node1: GraphNode = {
-        id: '1',
-        content: 'Node 1',
-        embedding: [1, 2, 3],
-      };
-      const node2: GraphNode = {
-        id: '2',
-        content: 'Node 2',
-        embedding: [4, 5, 6],
-      };
-      graph.addNode(node1);
-      graph.addNode(node2);
-      const edge: GraphEdge = {
-        source: '1',
-        target: '2',
-        weight: 0.5,
-        type: 'semantic',
-      };
-      graph.addEdge(edge);
-      expect(graph['edges'].length).toBe(2);
-    });
-  });
-
-  describe('createGraph', () => {
-    it("chunks and embeddings can't be empty", () => {
-      const graph = new GraphRAG(3);
-      const chunks: GraphChunk[] = [];
-      const embeddings: GraphEmbedding[] = [];
-      expect(() => graph.createGraph(chunks, embeddings)).toThrowError(
-        'Chunks and embeddings arrays must not be empty',
-      );
-    });
-    it('chunks and embeddings must have the same length', () => {
-      const graph = new GraphRAG(3);
-      const chunks: GraphChunk[] = [
-        {
-          text: 'Chunk 1',
-          metadata: {},
-        },
-        {
-          text: 'Chunk 2',
-          metadata: {},
-        },
-      ];
-      const embeddings: GraphEmbedding[] = [
-        {
-          vector: [1, 2, 3],
-        },
-      ];
-      expect(() => graph.createGraph(chunks, embeddings)).toThrowError(
-        'Chunks and embeddings must have the same length',
-      );
-    });
-    it('should return the top ranked nodes', () => {
-      const results = [
-        {
-          metadata: {
-            text: 'Chunk 1',
-          },
-          vector: [1, 2, 3],
-        },
-        {
-          metadata: {
-            text: 'Chunk 2',
-          },
-          vector: [4, 5, 6],
-        },
-        {
-          metadata: {
-            text: 'Chunk 3',
-          },
-          vector: [7, 8, 9],
-        },
-      ];
-
-      const chunks = results.map(result => ({
-        text: result?.metadata?.text,
-        metadata: result.metadata,
-      }));
-      const embeddings = results.map(result => ({
-        vector: result.vector,
-      }));
-
-      const graph = new GraphRAG(3);
-      graph.createGraph(chunks, embeddings);
-
-      const nodes = graph.getNodes();
-      expect(nodes.length).toBe(3);
-      expect(nodes[0]?.id).toBe('0');
-      expect(nodes[1]?.id).toBe('1');
-      expect(nodes[2]?.id).toBe('2');
-
-      const edges = graph.getEdges();
-      expect(edges.length).toBe(6);
-    });
-  });
-
-  describe('query', () => {
-    it("query embedding can't be empty", () => {
-      const graph = new GraphRAG(3);
-      const queryEmbedding: number[] = [];
-      expect(() => graph.query({ query: queryEmbedding, topK: 2, randomWalkSteps: 3, restartProb: 0.1 })).toThrowError(
-        `Query embedding must have dimension ${3}`,
-      );
-    });
-
-    it('topK must be greater than 0', () => {
-      const graph = new GraphRAG(3);
-      const queryEmbedding = [1, 2, 3];
-      const topK = 0;
-      expect(() => graph.query({ query: queryEmbedding, topK, randomWalkSteps: 3, restartProb: 0.1 })).toThrowError(
-        'TopK must be greater than 0',
-      );
-    });
-
-    it('randomWalkSteps must be greater than 0', () => {
-      const graph = new GraphRAG(3);
-      const queryEmbedding = [1, 2, 3];
-      const topK = 2;
-      const randomWalkSteps = 0;
-      expect(() => graph.query({ query: queryEmbedding, topK, randomWalkSteps, restartProb: 0.1 })).toThrowError(
-        'Random walk steps must be greater than 0',
-      );
-    });
-
-    it('restartProb must be between 0 and 1', () => {
-      const graph = new GraphRAG(3);
-      const queryEmbedding = [1, 2, 3];
-      const topK = 2;
-      const randomWalkSteps = 3;
-      const restartProb = -0.1;
-      expect(() => graph.query({ query: queryEmbedding, topK, randomWalkSteps, restartProb })).toThrowError(
-        'Restart probability must be between 0 and 1',
-      );
-    });
-
-    it('should return the top ranked nodes', () => {
-      const graph = new GraphRAG(3);
-      const node1: GraphNode = {
-        id: '1',
-        content: 'Node 1',
-        embedding: [1, 2, 3],
-      };
-      const node2: GraphNode = {
-        id: '2',
-        content: 'Node 2',
-        embedding: [11, 12, 13],
-      };
-      const node3: GraphNode = {
-        id: '3',
-        content: 'Node 3',
-        embedding: [21, 22, 23],
-      };
-      graph.addNode(node1);
-      graph.addNode(node2);
-      graph.addNode(node3);
-      graph.addEdge({
-        source: '1',
-        target: '2',
-        weight: 0.5,
-        type: 'semantic',
-      });
-      graph.addEdge({
-        source: '2',
-        target: '3',
-        weight: 0.7,
-        type: 'semantic',
-      });
-
-      const queryEmbedding = [15, 16, 17];
-      const topK = 2;
-      const randomWalkSteps = 3;
-      const restartProb = 0.1;
-      const rerankedResults = graph.query({ query: queryEmbedding, topK, randomWalkSteps, restartProb });
-
-      expect(rerankedResults.length).toBe(2);
-    });
-  });
-});