@mlx-node/trl 0.0.13 → 0.0.15

package/src/types.ts

@@ -0,0 +1,70 @@
+// Re-export types from @mlx-node/core
+export type { CompletionInfo, RewardOutput } from '@mlx-node/core';
+import type { RewardOutput } from '@mlx-node/core';
+
+// Re-export ChatRole from @mlx-node/lm (single source of truth)
+export type { ChatRole } from '@mlx-node/lm';
+
+export interface ChatMessage {
+  role: 'system' | 'user' | 'assistant' | 'tool' | (string & {});
+  content: string;
+}
+
+export interface CompletionMessage extends ChatMessage {}
+
+export type Completion = CompletionMessage[];
+
+export type DatasetSplit = 'train' | 'test' | (string & {});
+
+export interface DatasetExample {
+  prompt: ChatMessage[];
+  metadata?: Record<string, unknown>;
+}
+
+export interface XmlParseResult {
+  reasoning: string | null;
+  answer: string | null;
+  isStrictMatch: boolean;
+  isSoftMatch: boolean;
+  errors: string[];
+}
+
+export interface RewardComputationInput {
+  prompts: ChatMessage[][];
+  completions: Completion[];
+  answers: (string | null)[];
+}
+
+/**
+ * Unified reward function type for GRPO training.
+ *
+ * Takes an array of RewardOutput objects containing structured completion data.
+ * Returns rewards for each completion (one per output).
+ */
+export type RewardFunction<T = unknown> = (
+  outputs: RewardOutput[],
+  context: T,
+) => number[] | Float32Array | Promise<number[] | Float32Array>;
+
+export interface PromptFormatterOptions {
+  includeOneShot?: boolean;
+  oneShotExample?: {
+    question: string;
+    reasoning: string;
+    answer: string;
+  };
+}
+
+export type PromptTemplate = (question: string, options?: PromptFormatterOptions) => ChatMessage[];
+
+/**
+ * Converts a ChatMessage array to a string for reward function input
+ *
+ * This allows customization of how prompts are formatted as strings
+ * for different model architectures (Qwen3, Llama, etc.)
+ */
+export type PromptFormatter = (messages: ChatMessage[]) => string;
+
+export interface DatasetLoader {
+  load(split: DatasetSplit, limit?: number): Promise<DatasetExample[]>;
+}

package/src/utils/path-security.ts

@@ -0,0 +1,88 @@
+/**
+ * Path security utilities to prevent directory traversal attacks.
+ *
+ * These utilities ensure that user-provided paths stay within allowed directories,
+ * preventing malicious paths like `../../../etc/` from accessing arbitrary files.
+ */
+
+import { resolve as resolvePath, normalize, relative, isAbsolute } from 'node:path';
+
+/**
+ * Error thrown when a path traversal attempt is detected.
+ */
+export class PathTraversalError extends Error {
+  constructor(
+    public readonly resolvedPath: string,
+    public readonly allowedRoot: string,
+  ) {
+    super(`Path traversal detected: "${resolvedPath}" is outside allowed directory "${allowedRoot}"`);
+    this.name = 'PathTraversalError';
+  }
+}
+
+/**
+ * Validates that a resolved path is contained within an allowed root directory.
+ * Prevents path traversal attacks via '../' sequences.
+ *
+ * @param resolvedPath - The fully resolved absolute path to validate
+ * @param allowedRoot - The root directory that paths must be contained within
+ * @throws PathTraversalError if path escapes the allowed root
+ */
+export function validatePathContainment(resolvedPath: string, allowedRoot: string): void {
+  const normalizedPath = normalize(resolvedPath);
+  const normalizedRoot = normalize(allowedRoot);
+
+  // Get relative path from root to target
+  const relativePath = relative(normalizedRoot, normalizedPath);
+
+  // If relative path starts with '..' or is absolute, it's outside the root
+  if (relativePath.startsWith('..') || isAbsolute(relativePath)) {
+    throw new PathTraversalError(resolvedPath, allowedRoot);
+  }
+}
+
+/**
+ * Resolves a user-provided path and validates it stays within an allowed root.
+ *
+ * @param userPath - The user-provided path (may be relative or absolute)
+ * @param allowedRoot - The root directory that the path must be contained within
+ * @returns The resolved absolute path
+ * @throws PathTraversalError if the resolved path escapes the allowed root
+ */
+export function resolveAndValidatePath(userPath: string, allowedRoot: string): string {
+  const resolved = resolvePath(allowedRoot, userPath);
+  validatePathContainment(resolved, allowedRoot);
+  return resolved;
+}
+
+/**
+ * Options for configuring path validation behavior.
+ */
+export interface PathValidationOptions {
+  /**
+   * The root directory that all paths must be contained within.
+   * Defaults to process.cwd() if not specified.
+   */
+  allowedRoot?: string;
+}
+
+/**
+ * Get the allowed root directory from options or environment.
+ * Checks MLX_NODE_DATA_ROOT environment variable first, then falls back to cwd.
+ *
+ * @param options - Optional validation options
+ * @returns The allowed root directory path
+ */
+export function getAllowedRoot(options?: PathValidationOptions): string {
+  if (options?.allowedRoot) {
+    return options.allowedRoot;
+  }
+
+  // Check environment variable for configurable data root
+  const envRoot = process.env['MLX_NODE_DATA_ROOT'];
+  if (envRoot) {
+    return resolvePath(envRoot);
+  }
+
+  return process.cwd();
+}

package/src/utils/xml-parser.ts

@@ -0,0 +1,209 @@
+import { type XmlParseResult } from '../types.js';
+
+const REASONING_OPEN = '<reasoning>';
+const REASONING_CLOSE = '</reasoning>';
+const ANSWER_OPEN = '<answer>';
+const ANSWER_CLOSE = '</answer>';
+
+type ParserState =
+  | 'searchReasoningOpen'
+  | 'readReasoning'
+  | 'searchAnswerOpen'
+  | 'readAnswer'
+  | 'consumeTrailing'
+  | 'done';
+
+function isWhitespace(char: string): boolean {
+  return char === ' ' || char === '\n' || char === '\r' || char === '\t' || char === '\f' || char === '\v';
+}
+
+function sliceWithTag(text: string, openTag: string, closeTag: string): string | null {
+  const openIndex = text.indexOf(openTag);
+  if (openIndex === -1) return null;
+  const start = openIndex + openTag.length;
+  const closeIndex = text.indexOf(closeTag, start);
+  if (closeIndex === -1) return null;
+  return text.slice(start, closeIndex).trim();
+}
+
+export function parseXmlCot(content: string): XmlParseResult {
+  const text = content ?? '';
+  const errors: string[] = [];
+
+  let state: ParserState = 'searchReasoningOpen';
+  let index = 0;
+  const length = text.length;
+
+  let reasoning: string | null = null;
+  let answer: string | null = null;
+
+  let reasoningFound = false;
+  let answerFound = false;
+  let reasoningClosed = false;
+  let answerClosed = false;
+
+  let hasLeadingText = false;
+  let hasBetweenText = false;
+  let hasTrailingText = false;
+
+  let usedFallbackReasoning = false;
+  let usedFallbackAnswer = false;
+
+  while (state !== 'done') {
+    switch (state) {
+      case 'searchReasoningOpen': {
+        while (index < length) {
+          const char = text[index];
+          if (isWhitespace(char)) {
+            index += 1;
+            continue;
+          }
+          if (text.startsWith(REASONING_OPEN, index)) {
+            reasoningFound = true;
+            index += REASONING_OPEN.length;
+            state = 'readReasoning';
+            break;
+          }
+          hasLeadingText = true;
+          index += 1;
+        }
+        if (state === 'searchReasoningOpen') {
+          state = 'done';
+        }
+        break;
+      }
+      case 'readReasoning': {
+        const closingIndex = text.indexOf(REASONING_CLOSE, index);
+        if (closingIndex === -1) {
+          reasoning = text.slice(index).trim() || null;
+          errors.push('Unterminated <reasoning>...</reasoning> section.');
+          state = 'done';
+          break;
+        }
+        reasoning = text.slice(index, closingIndex).trim();
+        reasoningClosed = true;
+        index = closingIndex + REASONING_CLOSE.length;
+        state = 'searchAnswerOpen';
+        break;
+      }
+      case 'searchAnswerOpen': {
+        while (index < length) {
+          const char = text[index];
+          if (isWhitespace(char)) {
+            index += 1;
+            continue;
+          }
+          if (text.startsWith(ANSWER_OPEN, index)) {
+            answerFound = true;
+            index += ANSWER_OPEN.length;
+            state = 'readAnswer';
+            break;
+          }
+          hasBetweenText = true;
+          index += 1;
+        }
+        if (state === 'searchAnswerOpen') {
+          state = 'done';
+        }
+        break;
+      }
+      case 'readAnswer': {
+        const closingIndex = text.indexOf(ANSWER_CLOSE, index);
+        if (closingIndex === -1) {
+          answer = text.slice(index).trim() || null;
+          errors.push('Unterminated <answer>...</answer> section.');
+          state = 'done';
+          break;
+        }
+        answer = text.slice(index, closingIndex).trim();
+        answerClosed = true;
+        index = closingIndex + ANSWER_CLOSE.length;
+        state = 'consumeTrailing';
+        break;
+      }
+      case 'consumeTrailing': {
+        while (index < length) {
+          const char = text[index];
+          if (!isWhitespace(char)) {
+            hasTrailingText = true;
+          }
+          index += 1;
+        }
+        state = 'done';
+        break;
+      }
+    }
+  }
+
+  if (!reasoningFound) {
+    const fallback = sliceWithTag(text, REASONING_OPEN, REASONING_CLOSE);
+    if (fallback !== null) {
+      reasoning = fallback;
+      reasoningFound = true;
+      reasoningClosed = true;
+      usedFallbackReasoning = true;
+    }
+  }
+
+  if (!answerFound) {
+    const fallback = sliceWithTag(text, ANSWER_OPEN, ANSWER_CLOSE);
+    if (fallback !== null) {
+      answer = fallback;
+      answerFound = true;
+      answerClosed = true;
+      usedFallbackAnswer = true;
+    }
+  }
+
+  if (!reasoningFound) {
+    errors.push('Missing <reasoning>...</reasoning> section.');
+  }
+  if (!answerFound) {
+    errors.push('Missing <answer>...</answer> section.');
+  }
+
+  if (hasLeadingText) {
+    errors.push('XML format contains extra characters before <reasoning> section.');
+  }
+  if (hasBetweenText) {
+    errors.push('XML format contains extra characters between reasoning and answer sections.');
+  }
+  if (hasTrailingText) {
+    errors.push('XML format contains extra characters after </answer> section.');
+  }
+
+  const isSoftMatch = reasoningFound && answerFound && reasoningClosed && answerClosed;
+  const isStrictMatch =
+    isSoftMatch &&
+    !hasLeadingText &&
+    !hasBetweenText &&
+    !hasTrailingText &&
+    !usedFallbackReasoning &&
+    !usedFallbackAnswer;
+
+  return {
+    reasoning: reasoning ?? null,
+    answer: answer ?? null,
+    isStrictMatch,
+    isSoftMatch,
+    errors,
+  };
+}
+
+export function extractXmlAnswer(content: string): string | null {
+  const { answer } = parseXmlCot(content);
+  return answer;
+}
+
+export function extractXmlReasoning(content: string): string | null {
+  const { reasoning } = parseXmlCot(content);
+  return reasoning;
+}
+
+export function extractHashAnswer(text: string): string | null {
+  if (!text) return null;
+  const separator = text.indexOf('####');
+  if (separator === -1) return null;
+  const extracted = text.slice(separator + 4).trim();
+  return extracted.length ? extracted : null;
+}