mdast chunk (#1736)

* add naive chunking

* chunking

* refactor: enhance mdast stringify function to handle array input and improve chunking tests

* refactor: update mdastStringify to accept RootContent and improve type safety in tests

* docs: add example for chunking the mdast tree into sections

* test: add chunking tests for large markdown documents and improve coverage

* refactor: improve chunking logic in mdast function to handle nodes more efficiently

* refactor: update mdast function to accept Root type and adjust chunking logic in tests

Author: pelikhan

.vscode/settings.json

@@ -274,6 +274,7 @@
         "qwen",
         "RAAA",
         "redteam",
+        "remarkalerts",
         "remarkdetails",
         "resd",
         "resj",

docs/src/content/docs/reference/runtime/plugin-mdast.mdx

@@ -53,6 +53,17 @@ const updated = visit(root, `code`, (node) => {
 const markdown = await stringify(updated);
 ```
 
+- chunk the tree into sections
+
+```typescript
+const { parse, chunk } = await mdast();
+const root = parse(
+  "# Section 1\n\nContent 1\n\n## Subsection 1.1\n\nContent 1.1\n\n# Section 2\n\nContent 2",
+);
+const sections = chunk(root);
+console.log(sections);
+```
+
 In order to get type completion, you will need to install the `@types/mdast` package as a development dependency.
 
 ## Debugging trees

packages/core/src/anthropic.ts

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-import { ChatCompletionHandler, LanguageModel, ListModelsFunction } from "./chat.js";
+import type { ChatCompletionHandler, LanguageModel, ListModelsFunction } from "./chat.js";
 import {
   ANTHROPIC_MAX_TOKEN,
   MODEL_PROVIDER_ANTHROPIC,
@@ -13,7 +13,7 @@ import { approximateTokens } from "./tokens.js";
 import { resolveTokenEncoder } from "./encoders.js";
 import type Anthropic from "@anthropic-ai/sdk";
 import type AnthropicBedrock from "@anthropic-ai/bedrock-sdk";
-import {
+import type {
   ChatCompletionResponse,
   ChatCompletionToolCall,
   ChatCompletionUsage,
@@ -33,14 +33,15 @@ import {
 import { logError } from "./util.js";
 import { resolveUndiciProxyAgent } from "./proxy.js";
 import type { ProxyAgent } from "undici";
-import { MarkdownTrace } from "./trace.js";
-import { createFetch, FetchType } from "./fetch.js";
+import type { MarkdownTrace } from "./trace.js";
+import { createFetch } from "./fetch.js";
+import type { FetchType } from "./fetch.js";
 import { JSONLLMTryParse } from "./json5.js";
-import { LanguageModelConfiguration } from "./server/messages.js";
+import type { LanguageModelConfiguration } from "./server/messages.js";
 import { deleteUndefinedValues } from "./cleaners.js";
 import debug from "debug";
 import { providerFeatures } from "./features.js";
-import { LanguageModelInfo } from "./types.js";
+import type { LanguageModelInfo } from "./types.js";
 const dbg = debug("genaiscript:anthropic");
 const dbgMessages = debug("genaiscript:anthropic:msg");
 

packages/plugin-mdast/src/remarkdetails.ts

@@ -8,6 +8,12 @@ import { remark } from "remark";
 import { genaiscriptDebug } from "@genaiscript/core";
 const dbg = genaiscriptDebug("mdast:html:details");
 
+declare module "mdast" {
+  interface RootContentMap {
+    details: DetailsElement;
+  }
+}
+
 export interface RemarkDetailsOptions {}
 
 export interface DetailsElement extends Parent {

packages/plugin-mdast/src/unified.ts

@@ -1,13 +1,15 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-import type { Root } from "mdast";
+import type { Root, RootContent } from "mdast";
 import type { WorkspaceFile } from "@genaiscript/core";
 import { checkRuntime, filenameOrFileToContent, genaiscriptDebug } from "@genaiscript/core";
 import type { Processor } from "unified";
 import remarkGitHubAlerts from "./remarkalerts.js";
 import type { GitHubAlertMarker } from "./remarkalerts.js";
-import remarkDetails, { DetailsElement, SummaryElement } from "./remarkdetails.js";
+import remarkDetails from "./remarkdetails.js";
+import type { DetailsElement, SummaryElement } from "./remarkdetails.js";
+import { approximateTokens } from "@genaiscript/core";
 const dbg = genaiscriptDebug("mdast");
 
 export interface MdAstOptions {
@@ -64,14 +66,15 @@ export async function mdast(options?: MdAstOptions) {
     return processed as Root;
   };
 
-  const mdastStringify = (root: Root, options?: {}): string => {
+  const mdastStringify = (root: Root | RootContent[], stringifyOptions?: object): string => {
     if (!root) return "";
 
     dbg(`stringify`);
     const processor = unified();
     usePlugins(processor, "stringify");
+    // @ts-expect-error - TypeScript doesn't recognize the handlers option
     processor.use(stringify, {
-      ...(options || {}),
+      ...(stringifyOptions || {}),
       handlers: {
         githubAlertMarker(node: GitHubAlertMarker) {
           return node.value;
@@ -83,15 +86,101 @@ export async function mdast(options?: MdAstOptions) {
           return `<summary>${node.children.map((child) => processor.stringify(child)).join("")}</summary>`;
         },
       },
-    } as any);
+    });
 
-    const result = processor.stringify(root);
+    const n = Array.isArray(root) ? ({ type: "root", children: root } satisfies Root) : root;
+    const result = processor.stringify(n);
     return String(result);
   };
 
+  const mdChunk = (
+    nodes: Root | RootContent[],
+    maxTokens: number,
+    chunkOptions?: {
+      tokenize: (text: string) => number;
+    },
+  ): RootContent[][] => {
+    const { tokenize = approximateTokens } = chunkOptions || {};
+    if (!nodes) return [];
+    if (!Array.isArray(nodes)) {
+      if (nodes.type !== "root") throw new Error("Expected nodes to be an array or a Root type");
+      nodes = nodes.children || [];
+    }
+    if (nodes.length === 0) return [];
+
+    const chunks: RootContent[][] = [];
+    let currentChunk: RootContent[] = [];
+    let currentTokenCount = 0;
+
+    const measure = (ns: RootContent[]): number => tokenize(mdastStringify(ns));
+
+    // Process nodes in order, never reordering them
+    for (let i = 0; i < nodes.length; i++) {
+      const node = nodes[i];
+      const nodeTokens = measure([node]);
+
+      // If adding this node would exceed the limit and we have content in current chunk
+      if (currentTokenCount + nodeTokens > maxTokens && currentChunk.length > 0) {
+        // For headings, try to keep them with their content by looking ahead
+        if (node.type === "heading") {
+          // Look ahead to see how much content follows this heading
+          let headingContentSize = nodeTokens;
+          let nextHeadingIndex = i + 1;
+
+          // Find content that belongs to this heading (until next heading of same or higher level)
+          while (nextHeadingIndex < nodes.length) {
+            const nextNode = nodes[nextHeadingIndex];
+            if (nextNode.type === "heading" && nextNode.depth <= node.depth) {
+              break; // Found a heading of same or higher level
+            }
+            headingContentSize += measure([nextNode]);
+            nextHeadingIndex++;
+          }
+
+          // If the heading + its content can fit in a new chunk, start a new chunk
+          if (headingContentSize <= maxTokens) {
+            chunks.push(currentChunk);
+            currentChunk = [];
+            currentTokenCount = 0;
+          }
+          // Otherwise, just finalize current chunk and continue
+          else if (currentChunk.length > 0) {
+            chunks.push(currentChunk);
+            currentChunk = [];
+            currentTokenCount = 0;
+          }
+        } else {
+          // For non-heading nodes, just start a new chunk
+          chunks.push(currentChunk);
+          currentChunk = [];
+          currentTokenCount = 0;
+        }
+      }
+
+      // Add the current node to the chunk
+      currentChunk.push(node);
+      currentTokenCount += nodeTokens;
+
+      // If this single node exceeds maxTokens, put it in its own chunk
+      if (nodeTokens > maxTokens && currentChunk.length === 1) {
+        chunks.push(currentChunk);
+        currentChunk = [];
+        currentTokenCount = 0;
+      }
+    }
+
+    // Add final chunk if it has content
+    if (currentChunk.length > 0) {
+      chunks.push(currentChunk);
+    }
+
+    return chunks;
+  };
+
   return Object.freeze({
     parse: mdastParse,
     stringify: mdastStringify,
+    chunk: mdChunk,
     visit,
     visitParents,
     inspect,