Implement custom chunk strategy for yaml text

Author: JHZhang2736

src/chunker.ts

@@ -0,0 +1,65 @@
+import {
+    ChunkResult,
+} from "@magda/semantic-indexer-sdk";
+import { RecursiveCharacterTextSplitter } from "@langchain/textsplitters";
+
+export class Chunker {
+    private chunkSize: number;
+    private overlap: number;
+    private splitter: RecursiveCharacterTextSplitter;
+
+    constructor(chunkSize: number, overlap: number) {
+        if (overlap >= chunkSize) {
+            throw new Error("Overlap must be smaller than chunk size");
+        }
+        this.chunkSize = chunkSize * 4;
+        this.overlap = overlap * 4;
+        this.splitter = new RecursiveCharacterTextSplitter({
+            separators: [
+                "\nColumns",
+                "\nDescription",
+                "\nTitle",
+                "\nKeywords",
+                "\nThemes",
+                "\nTemporal coverage",
+                "\nLanguages",
+                "\nPublisher",
+                "\n",
+                " ",
+                ""
+            ],
+            chunkSize: this.chunkSize,
+            chunkOverlap: this.overlap,
+            keepSeparator: true,
+            stripWhitespace: false,
+        });
+    }
+
+    async chunk(text: string): Promise<ChunkResult[]> {
+        if (!text) {
+            return [];
+        }
+
+        if (text.length <= this.chunkSize) {
+            return [{
+                text: text,
+                position: 0,
+                length: text.length,
+                overlap: 0
+            }];
+        }
+
+        const chunks = await this.splitter.splitTextWithMetadata(text);
+
+        for (let i = 1; i < chunks.length; i++) {
+            if (chunks[i].overlap === 0 && chunks[i].text.trim() === '') {
+                chunks[i-1].text = chunks[i-1].text + chunks[i].text;
+                chunks[i-1].length = chunks[i-1].length + chunks[i].length;
+                chunks.splice(i, 1);
+                i--;
+            }
+        }
+
+        return chunks;
+    }
+}

src/createEmbeddingText.ts

@@ -109,14 +109,14 @@ export const createEmbeddingText: CreateEmbeddingText = async ({
     const yamlText = toYaml({
         Title: title || undefined,
         Format: format || "CSV",
-        Description: description || undefined,
         "File name": fileName || undefined,
         "Temporal coverage": temporalCoverage || undefined,
         Licence: licence || undefined,
         Publisher: publisher || undefined,
         Themes: themes || undefined,
         Keywords: keywords || undefined,
         Languages: languages || undefined,
+        Description: description || undefined,
         Columns: columns || undefined
     });
 

src/index.ts

@@ -1,20 +1,29 @@
 import semanticIndexer, {
     SemanticIndexerOptions,
     commonYargs,
+    ChunkStrategyType,
 } from "@magda/semantic-indexer-sdk";
 import { csvSemanticIndexerArgs } from "./csvSemanticIndexerArgs.js";
 import { createEmbeddingText } from "./createEmbeddingText.js";
+import { Chunker } from "./chunker.js";
 
 const port = csvSemanticIndexerArgs.port;
 const args = commonYargs(port, `http://localhost:${port}`);
 
+const chunker = new Chunker(csvSemanticIndexerArgs.chunkSizeLimit, csvSemanticIndexerArgs.overlap);
+
+const chunkStrategy: ChunkStrategyType = async (text: string) => {
+    return await chunker.chunk(text);
+};
+
 const options: SemanticIndexerOptions = {
     argv: args,
     id: csvSemanticIndexerArgs.id,
     itemType: "storageObject",
     formatTypes: ["csv"],
     autoDownloadFile: true,
-    createEmbeddingText: createEmbeddingText
+    createEmbeddingText: createEmbeddingText,
+    chunkStrategy: chunkStrategy,
 };
 
 semanticIndexer(options);

src/test/chunker.spec.ts

@@ -0,0 +1,87 @@
+import { expect } from "chai";
+import { Chunker } from "../chunker.js";
+
+describe("Chunker", () => {
+    it("should chunk yaml text with proper overlapping", async () => {
+        const text = `Title: Dataset Title
+Publisher: Example Publisher
+Themes: [Environment, Climate]
+Keywords: [weather, rainfall]
+Languages: [en, zh]
+Description: This is a long description for testing chunker logic
+Columns:
+  - column1
+  - column2
+  - column3
+  - column4
+  - column5
+`;
+        const chunker = new Chunker(10, 3);
+        const chunks = await chunker.chunk(text);
+
+        expect(chunks).to.not.be.empty;
+        
+        chunks.forEach((chunk) => {
+            expect(chunk.length).to.equal(chunk.text.length);
+            expect(text.slice(chunk.position, chunk.position + chunk.length))
+                .to.equal(chunk.text);
+        });
+
+        const reconstructed = chunks.map(chunk => 
+            chunk.text.slice(chunk.overlap)
+        ).join('');
+        expect(reconstructed).to.equal(text);
+    });
+
+    it("should chunk yaml text with newline characters in description", async () => {
+        const text = `
+Title: Another Dataset
+Description: This is a very long description\nthat will be chunked into\nmultiple chunks.
+Publisher: Test Publisher
+Themes: [Data, Analysis]
+Columns:
+  - column1
+  - column2
+  - column3
+  - column4
+  - column5
+`;
+        const chunker = new Chunker(40, 8);
+        const chunks = await chunker.chunk(text);
+
+        expect(chunks).to.not.be.empty;
+        chunks.forEach((chunk) => {
+            expect(chunk.length).to.equal(chunk.text.length);
+            expect(text.slice(chunk.position, chunk.position + chunk.length))
+                .to.equal(chunk.text);
+        });
+
+        const reconstructed = chunks.map(chunk => 
+            chunk.text.slice(chunk.overlap)
+        ).join('');
+        expect(reconstructed).to.equal(text);
+    });
+
+    it("should handle text smaller than chunk size", async () => {
+        const text = `Title: Small Dataset`;
+        const chunker = new Chunker(50, 5);
+        const chunks = await chunker.chunk(text);
+        
+        expect(chunks).to.deep.equal([{
+            text,
+            length: text.length,
+            position: 0,
+            overlap: 0
+        }]);
+    });
+
+    it("should throw error when overlap >= chunk size", () => {
+        expect(() => new Chunker(20, 25)).to.throw("Overlap must be smaller than chunk size");
+    });
+
+    it("should handle empty text", async () => {
+        const chunker = new Chunker(50, 5);
+        const chunks = await chunker.chunk("");
+        expect(chunks).to.be.empty;
+    });
+});