@caleblawson/rag/src/tools/document-chunker.ts

The Retrieval-Augmented Generation (RAG) module contains document processing and embedding utilities.

import { createTool } from '@mastra/core/tools';
import { z } from 'zod';

import type { MDocument, ChunkParams } from '../document';

export const createDocumentChunkerTool = ({
  doc,
  params = {
    strategy: 'recursive',
    size: 512,
    overlap: 50,
    separator: '\n',
  },
}: {
  doc: MDocument;
  params?: ChunkParams;
}): ReturnType<typeof createTool> => {
  return createTool({
    id: `Document Chunker ${params.strategy} ${params.size}`,
    inputSchema: z.object({}),
    description: `Chunks document using ${params.strategy} strategy with size ${params.size} and ${params.overlap} overlap`,
    execute: async () => {
      const chunks = await doc.chunk(params);

      return {
        chunks,
      };
    },
  });
};