create-walle 0.9.25 → 0.9.26

package/template/wall-e/llm/mlx.js

@@ -7,6 +7,32 @@ const fs = require('fs');
 
 const MLX_VENV = path.join(os.homedir(), '.mlx-env');
 const MLX_PYTHON = path.join(MLX_VENV, 'bin', 'python3');
+const WALLE_LOCAL_MODEL_ROOT = path.join(os.homedir(), '.walle', 'models');
+const GEMMA_E4B_QAT_REPO = 'mlx-community/gemma-4-E4B-it-qat-4bit';
+const GEMMA_E4B_TOOLUSE_MODEL_ID = 'gemma4-e4b-it-qat-tooluse-v1';
+const GEMMA_E4B_TOOLUSE_MODEL_DIR = path.join(WALLE_LOCAL_MODEL_ROOT, GEMMA_E4B_TOOLUSE_MODEL_ID);
+const GEMMA_E4B_TOOLUSE_ADAPTER_PATH = path.join(GEMMA_E4B_TOOLUSE_MODEL_DIR, 'adapter');
+const GEMMA_E4B_TOOLUSE_METADATA_PATH = path.join(GEMMA_E4B_TOOLUSE_MODEL_DIR, 'candidate.json');
+const GEMMA_E4B_TOOLUSE_V2_MODEL_ID = 'gemma4-e4b-it-qat-tooluse-v2';
+const GEMMA_E4B_TOOLUSE_V2_MODEL_DIR = path.join(WALLE_LOCAL_MODEL_ROOT, GEMMA_E4B_TOOLUSE_V2_MODEL_ID);
+const GEMMA_E4B_TOOLUSE_V2_ADAPTER_PATH = path.join(GEMMA_E4B_TOOLUSE_V2_MODEL_DIR, 'adapter');
+const GEMMA_E4B_TOOLUSE_V2_METADATA_PATH = path.join(GEMMA_E4B_TOOLUSE_V2_MODEL_DIR, 'candidate.json');
+const GEMMA_E4B_TOOLUSE_V3_MODEL_ID = 'gemma4-e4b-it-qat-tooluse-v3';
+const GEMMA_E4B_TOOLUSE_V3_MODEL_DIR = path.join(WALLE_LOCAL_MODEL_ROOT, GEMMA_E4B_TOOLUSE_V3_MODEL_ID);
+const GEMMA_E4B_TOOLUSE_V3_ADAPTER_PATH = path.join(GEMMA_E4B_TOOLUSE_V3_MODEL_DIR, 'adapter');
+const GEMMA_E4B_TOOLUSE_V3_METADATA_PATH = path.join(GEMMA_E4B_TOOLUSE_V3_MODEL_DIR, 'candidate.json');
+const GEMMA_E4B_TOOLUSE_V4_MODEL_ID = 'gemma4-e4b-it-qat-tooluse-v4';
+const GEMMA_E4B_TOOLUSE_V4_MODEL_DIR = path.join(WALLE_LOCAL_MODEL_ROOT, GEMMA_E4B_TOOLUSE_V4_MODEL_ID);
+const GEMMA_E4B_TOOLUSE_V4_ADAPTER_PATH = path.join(GEMMA_E4B_TOOLUSE_V4_MODEL_DIR, 'adapter');
+const GEMMA_E4B_TOOLUSE_V4_METADATA_PATH = path.join(GEMMA_E4B_TOOLUSE_V4_MODEL_DIR, 'candidate.json');
+const GEMMA_E4B_TOOLUSE_V5_MODEL_ID = 'gemma4-e4b-it-qat-tooluse-v5';
+const GEMMA_E4B_TOOLUSE_V5_MODEL_DIR = path.join(WALLE_LOCAL_MODEL_ROOT, GEMMA_E4B_TOOLUSE_V5_MODEL_ID);
+const GEMMA_E4B_TOOLUSE_V5_ADAPTER_PATH = path.join(GEMMA_E4B_TOOLUSE_V5_MODEL_DIR, 'adapter');
+const GEMMA_E4B_TOOLUSE_V5_METADATA_PATH = path.join(GEMMA_E4B_TOOLUSE_V5_MODEL_DIR, 'candidate.json');
+const GEMMA_E4B_TOOLUSE_V6_MODEL_ID = 'gemma4-e4b-it-qat-tooluse-v6';
+const GEMMA_E4B_TOOLUSE_V6_MODEL_DIR = path.join(WALLE_LOCAL_MODEL_ROOT, GEMMA_E4B_TOOLUSE_V6_MODEL_ID);
+const GEMMA_E4B_TOOLUSE_V6_ADAPTER_PATH = path.join(GEMMA_E4B_TOOLUSE_V6_MODEL_DIR, 'adapter');
+const GEMMA_E4B_TOOLUSE_V6_METADATA_PATH = path.join(GEMMA_E4B_TOOLUSE_V6_MODEL_DIR, 'candidate.json');
 
 // Map of known model aliases to HuggingFace MLX repos
 const MLX_MODEL_REGISTRY = {
@@ -17,6 +43,38 @@ const MLX_MODEL_REGISTRY = {
   'phi4':            'mlx-community/phi-4-4bit',
   'phi4:latest':     'mlx-community/phi-4-4bit',
   'gemma3:4b':       'mlx-community/gemma-3-4b-it-4bit',
+  'gemma4:e4b':       GEMMA_E4B_QAT_REPO,
+  'gemma4:e4b-it-qat': GEMMA_E4B_QAT_REPO,
+  'gemma4:e4b-it-qat-tooluse-v1': {
+    repo: GEMMA_E4B_QAT_REPO,
+    adapterPath: GEMMA_E4B_TOOLUSE_ADAPTER_PATH,
+    metadataPath: GEMMA_E4B_TOOLUSE_METADATA_PATH,
+  },
+  'gemma4:e4b-it-qat-tooluse-v2': {
+    repo: GEMMA_E4B_QAT_REPO,
+    adapterPath: GEMMA_E4B_TOOLUSE_V2_ADAPTER_PATH,
+    metadataPath: GEMMA_E4B_TOOLUSE_V2_METADATA_PATH,
+  },
+  'gemma4:e4b-it-qat-tooluse-v3': {
+    repo: GEMMA_E4B_QAT_REPO,
+    adapterPath: GEMMA_E4B_TOOLUSE_V3_ADAPTER_PATH,
+    metadataPath: GEMMA_E4B_TOOLUSE_V3_METADATA_PATH,
+  },
+  'gemma4:e4b-it-qat-tooluse-v4': {
+    repo: GEMMA_E4B_QAT_REPO,
+    adapterPath: GEMMA_E4B_TOOLUSE_V4_ADAPTER_PATH,
+    metadataPath: GEMMA_E4B_TOOLUSE_V4_METADATA_PATH,
+  },
+  'gemma4:e4b-it-qat-tooluse-v5': {
+    repo: GEMMA_E4B_QAT_REPO,
+    adapterPath: GEMMA_E4B_TOOLUSE_V5_ADAPTER_PATH,
+    metadataPath: GEMMA_E4B_TOOLUSE_V5_METADATA_PATH,
+  },
+  'gemma4:e4b-it-qat-tooluse-v6': {
+    repo: GEMMA_E4B_QAT_REPO,
+    adapterPath: GEMMA_E4B_TOOLUSE_V6_ADAPTER_PATH,
+    metadataPath: GEMMA_E4B_TOOLUSE_V6_METADATA_PATH,
+  },
   'llama3.2:3b':     'mlx-community/Llama-3.2-3B-Instruct-4bit',
   'mistral:7b':      'mlx-community/Mistral-7B-Instruct-v0.3-4bit',
 };
@@ -33,18 +91,392 @@ function execAsync(cmd, args, opts = {}) {
   });
 }
 
+function resolveInferenceTimeoutMs(config = {}) {
+  const raw = config.inferenceTimeoutMs ?? process.env.MLX_INFERENCE_TIMEOUT_MS;
+  if (raw === undefined || raw === null || raw === '') return 300000;
+  const parsed = Number(raw);
+  return Number.isFinite(parsed) && parsed > 0 ? parsed : 300000;
+}
+
+function buildInferenceExecOptions(config = {}, signal) {
+  return {
+    timeout: resolveInferenceTimeoutMs(config),
+    killSignal: 'SIGTERM',
+    signal,
+    env: { ...process.env, TOKENIZERS_PARALLELISM: 'false' },
+  };
+}
+
+function normalizeRegistryEntry(entry) {
+  if (!entry) return null;
+  if (typeof entry === 'string') return { repo: entry, adapterPath: null, metadataPath: null };
+  if (typeof entry === 'object' && entry.repo) {
+    return {
+      repo: entry.repo,
+      adapterPath: entry.adapterPath || null,
+      metadataPath: entry.metadataPath || null,
+    };
+  }
+  return null;
+}
+
 /**
- * Resolve a model name to an MLX HuggingFace repo.
+ * Resolve a model name to an MLX HuggingFace repo plus optional adapter.
  * Accepts: registry aliases, full HF repo names (org/model), or mlx-community/ prefixed names.
  */
-function resolveModel(model) {
+function resolveModelSpec(model, config = {}) {
   if (!model) return null;
-  if (MLX_MODEL_REGISTRY[model]) return MLX_MODEL_REGISTRY[model];
+  const requested = String(model).trim();
+  const registered = normalizeRegistryEntry(MLX_MODEL_REGISTRY[requested]);
+  if (registered) {
+    return {
+      alias: requested,
+      repo: registered.repo,
+      adapterPath: config.adapterPath || registered.adapterPath || null,
+      metadataPath: registered.metadataPath || null,
+    };
+  }
   // Validate model name to prevent arbitrary HF repo downloads
-  if (!/^[\w\-.:\/]+$/.test(model)) return null;
-  if (model.includes('/')) return model; // Full HF repo
+  if (!/^[\w\-.:\/]+$/.test(requested)) return null;
+  if (requested.includes('/')) {
+    return {
+      alias: null,
+      repo: requested,
+      adapterPath: config.adapterPath || null,
+      metadataPath: null,
+    };
+  }
   // Try with mlx-community prefix
-  return `mlx-community/${model}`;
+  return {
+    alias: requested,
+    repo: `mlx-community/${requested}`,
+    adapterPath: config.adapterPath || null,
+    metadataPath: null,
+  };
+}
+
+function resolveModel(model) {
+  return resolveModelSpec(model)?.repo || null;
+}
+
+function resolveMlxPython(config = {}) {
+  return config.python || process.env.MLX_PYTHON || MLX_PYTHON;
+}
+
+function normalizeToolSpec(tool = {}) {
+  const fn = tool.function || {};
+  const schema = tool.input_schema || tool.inputSchema || tool.parameters || fn.parameters || {};
+  return {
+    name: tool.name || fn.name || '',
+    description: tool.description || fn.description || '',
+    schema: schema && typeof schema === 'object' ? schema : {},
+  };
+}
+
+function normalizeToolCallFormat(format) {
+  return format === 'bare-json' ? 'bare-json' : 'fenced';
+}
+
+function buildToolContext(tools = [], opts = {}) {
+  const normalized = tools.map(normalizeToolSpec).filter((tool) => tool.name);
+  if (normalized.length === 0) return '';
+  const toolCallFormat = normalizeToolCallFormat(opts.toolCallFormat || opts.format);
+  const instruction = toolCallFormat === 'bare-json'
+    ? '\n\nYou have access to these tools. To use a tool, respond with exactly one JSON object and nothing else. The first non-whitespace character must be "{". Do not use markdown fences. Do not output <|channel>thought.\nSchema: {"tool":"tool_name","input":{"arg":"value"}}\n\nAvailable tools:\n'
+    : '\n\nYou have access to these tools. To use a tool, respond with a JSON block:\n' +
+      '```json\n{"tool": "tool_name", "input": {"arg": "value"}}\n```\n\nAvailable tools:\n';
+  return instruction +
+    normalized.map((tool) => {
+      const schema = Object.keys(tool.schema || {}).length ? tool.schema : { type: 'object', properties: {} };
+      return `- **${tool.name}**: ${tool.description || 'No description'}\n  Parameters: ${JSON.stringify(schema)}`;
+    }).join('\n');
+}
+
+function toolCallInput(call = {}) {
+  const fn = call.function || {};
+  return normalizeToolInput(fn.arguments ?? call.input ?? call.arguments ?? {});
+}
+
+function renderToolCallJsonBlock(call = {}, opts = {}) {
+  const fn = call.function || {};
+  const name = fn.name || call.name || '';
+  const payload = JSON.stringify({ tool: name, input: toolCallInput(call) });
+  return normalizeToolCallFormat(opts.toolCallFormat || opts.format) === 'bare-json'
+    ? payload
+    : '```json\n' + payload + '\n```';
+}
+
+function messageTextContent(message = {}) {
+  if (typeof message.content === 'string') return message.content;
+  if (!Array.isArray(message.content)) return '';
+  return message.content
+    .filter((block) => block?.type === 'text')
+    .map((block) => block.text || '')
+    .filter(Boolean)
+    .join('\n');
+}
+
+function renderToolResultMessage(message = {}) {
+  const id = message.tool_call_id || message.toolCallId || message.tool_use_id || 'tool';
+  const content = typeof message.content === 'string'
+    ? message.content
+    : Array.isArray(message.content)
+      ? message.content.map((block) => block?.text || block?.content || '').filter(Boolean).join('\n')
+      : JSON.stringify(message.content ?? '');
+  return {
+    role: 'user',
+    content: `[Tool result for ${id}]: ${content}`,
+  };
+}
+
+function renderMlxProviderMessages(message = {}, opts = {}) {
+  if (!message || typeof message !== 'object') return [];
+
+  if (message.role === 'tool') {
+    return [renderToolResultMessage(message)];
+  }
+
+  if (message.role === 'assistant' && Array.isArray(message.tool_calls) && message.tool_calls.length) {
+    const parts = [];
+    const text = messageTextContent(message).trim();
+    if (text) parts.push(text);
+    for (const call of message.tool_calls) {
+      parts.push(renderToolCallJsonBlock(call, opts));
+    }
+    return [{ role: 'assistant', content: parts.join('\n\n') }];
+  }
+
+  if (typeof message.content === 'string') {
+    return [{ role: message.role, content: message.content }];
+  }
+
+  if (Array.isArray(message.content)) {
+    const out = [];
+    const text = messageTextContent(message);
+    if (text) out.push({ role: message.role, content: text });
+    const toolResults = message.content.filter((block) => block?.type === 'tool_result');
+    for (const tr of toolResults) {
+      const resultText = typeof tr.content === 'string' ? tr.content
+        : Array.isArray(tr.content) ? tr.content.map((c) => c.text || '').join('\n')
+        : JSON.stringify(tr.content);
+      out.push({ role: 'user', content: `[Tool result for ${tr.tool_use_id}]: ${resultText}` });
+    }
+    return out;
+  }
+
+  return [];
+}
+
+function prepareMessagesForMlxProvider({
+  messages = [],
+  tools = [],
+  system = null,
+  compactDuplicates = true,
+  toolCallFormat = 'fenced',
+} = {}) {
+  const chatMessages = [];
+  if (system) chatMessages.push({ role: 'system', content: system });
+  for (const message of messages || []) {
+    chatMessages.push(...renderMlxProviderMessages(message, { toolCallFormat }));
+  }
+
+  const toolContext = buildToolContext(tools || [], { toolCallFormat });
+  if (toolContext) {
+    if (chatMessages.length > 0 && chatMessages[0].role === 'system') {
+      chatMessages[0] = {
+        ...chatMessages[0],
+        content: `${chatMessages[0].content || ''}${toolContext}`,
+      };
+    } else {
+      chatMessages.unshift({ role: 'system', content: toolContext.trim() });
+    }
+  }
+
+  return compactDuplicates ? compactDuplicateToolResults(chatMessages) : chatMessages;
+}
+
+function duplicateToolResultSignature(content, { minLength = 1200 } = {}) {
+  const text = String(content || '');
+  const match = text.match(/^\[Tool result for [^\]]+\]:\s*([\s\S]*)$/);
+  if (!match) return null;
+  const body = match[1].trim();
+  if (body.length < minLength) return null;
+  return body;
+}
+
+function compactDuplicateToolResults(messages = [], opts = {}) {
+  const seen = new Set();
+  const duplicateIndexes = new Set();
+  for (let i = messages.length - 1; i >= 0; i -= 1) {
+    const message = messages[i] || {};
+    if (message.role !== 'user') continue;
+    const signature = duplicateToolResultSignature(message.content, opts);
+    if (!signature) continue;
+    if (seen.has(signature)) duplicateIndexes.add(i);
+    else seen.add(signature);
+  }
+  if (duplicateIndexes.size === 0) return messages;
+  return messages.map((message, idx) => {
+    if (!duplicateIndexes.has(idx)) return message;
+    const id = String(message.content || '').match(/^\[Tool result for ([^\]]+)\]/)?.[1] || 'tool';
+    return {
+      ...message,
+      content: `[Tool result for ${id}]: Duplicate large tool result omitted; the latest identical result remains later in this prompt. Use that latest visible result instead of repeating the same read.`,
+    };
+  });
+}
+
+function extractJsonObject(text) {
+  const value = String(text || '');
+  const fenced = value.match(/```(?:json)?\s*\n?(\{[\s\S]*?\})\s*\n?```/i);
+  if (fenced) return fenced[1];
+
+  const start = value.indexOf('{');
+  if (start < 0) return null;
+
+  let depth = 0;
+  let inString = false;
+  let escaped = false;
+  for (let i = start; i < value.length; i += 1) {
+    const ch = value[i];
+    if (inString) {
+      if (escaped) escaped = false;
+      else if (ch === '\\') escaped = true;
+      else if (ch === '"') inString = false;
+      continue;
+    }
+    if (ch === '"') {
+      inString = true;
+    } else if (ch === '{') {
+      depth += 1;
+    } else if (ch === '}') {
+      depth -= 1;
+      if (depth === 0) return value.slice(start, i + 1);
+    }
+  }
+  return null;
+}
+
+function normalizeToolInput(input) {
+  if (!input) return {};
+  if (typeof input === 'string') {
+    try {
+      const parsed = JSON.parse(input);
+      return parsed && typeof parsed === 'object' ? parsed : {};
+    } catch {
+      return {};
+    }
+  }
+  return typeof input === 'object' ? input : {};
+}
+
+function relativizeBenchPath(value) {
+  if (typeof value !== 'string') return value;
+  const match = value.match(/(?:^|\/)bench-(?!sandbox(?:\/|$))[^/]+\/(.+)$/);
+  return match ? match[1] : value;
+}
+
+function normalizeGeneratedToolCall(name, input = {}) {
+  let normalizedName = String(name || '');
+  if (/^(ls|ls_directory)$/i.test(normalizedName)) normalizedName = 'list_directory';
+  if (/^glob_files$/i.test(normalizedName)) normalizedName = 'glob';
+  const normalized = { ...normalizeToolInput(input) };
+  for (const key of ['file_path', 'path', 'directory', 'cwd']) {
+    if (normalized[key]) normalized[key] = relativizeBenchPath(normalized[key]);
+  }
+  if ((normalizedName === 'read_file' || normalizedName === 'write_file' || normalizedName === 'edit_file') && normalized.path && !normalized.file_path) {
+    normalized.file_path = normalized.path;
+    delete normalized.path;
+  }
+  if (normalizedName === 'list_directory' && normalized.path && !normalized.directory) {
+    normalized.directory = normalized.path;
+    delete normalized.path;
+  }
+  if (normalizedName === 'run_shell') delete normalized.cwd;
+  return { name: normalizedName, input: normalized };
+}
+
+function parseLooseKeyValues(text) {
+  const out = {};
+  const pairRe = /"([A-Za-z_][A-Za-z0-9_]*)"\s*:\s*("(?:\\.|[^"\\])*"|-?\d+(?:\.\d+)?|true|false|null)/g;
+  let match;
+  while ((match = pairRe.exec(String(text || '')))) {
+    try {
+      out[match[1]] = JSON.parse(match[2]);
+    } catch {}
+  }
+  return out;
+}
+
+function recoverGemmaToolCallSegment(segment) {
+  const text = String(segment || '');
+  const named = text.match(/^\s*call:([^\s{<]+)/i);
+  const toolObjectName = text.match(/"tool"\s*:\s*"([^"]+)"/i) || text.match(/"name"\s*:\s*"([^"]+)"/i);
+  const name = named?.[1] || toolObjectName?.[1] || null;
+  if (!name) return null;
+
+  const argsStart = text.search(/\{+(?="?(?:command|file_path|path|directory|pattern|file_pattern|patch_text|old_string|new_string|content|limit|offset)"?\s*:)/);
+  if (argsStart >= 0) {
+    const argsCandidate = text.slice(argsStart).replace(/^\{+/, '{');
+    const argsText = extractJsonObject(argsCandidate);
+    if (argsText) {
+      try {
+        return normalizeGeneratedToolCall(name, JSON.parse(argsText));
+      } catch {}
+    }
+  }
+
+  const directArgs = named ? extractJsonObject(text.slice(named.index + named[0].length)) : null;
+  if (directArgs) {
+    try {
+      return normalizeGeneratedToolCall(name, JSON.parse(directArgs));
+    } catch {}
+  }
+
+  const loose = parseLooseKeyValues(text);
+  delete loose.tool;
+  delete loose.name;
+  return Object.keys(loose).length ? normalizeGeneratedToolCall(name, loose) : null;
+}
+
+function recoverMalformedBareJsonToolCall(text) {
+  const value = String(text || '');
+  const toolObjectName = value.match(/"tool"\s*:\s*"([^"]+)"/i) || value.match(/"name"\s*:\s*"([^"]+)"/i);
+  const name = toolObjectName?.[1] || null;
+  if (!name) return null;
+
+  const inputStart = value.search(/"input"\s*:\s*\{/i);
+  const inputText = inputStart >= 0 ? value.slice(inputStart) : value;
+  const loose = parseLooseKeyValues(inputText);
+  delete loose.tool;
+  delete loose.name;
+  return normalizeGeneratedToolCall(name, loose);
+}
+
+function parseToolCallText(text) {
+  const value = String(text || '');
+  const gemmaSegments = value.split(/<\|tool_call\>/i).slice(1);
+  for (const segment of gemmaSegments) {
+    const toolSegment = segment.split(/<tool_call\|>/i)[0];
+    const recovered = recoverGemmaToolCallSegment(toolSegment);
+    if (recovered?.name) return recovered;
+  }
+
+  const jsonText = extractJsonObject(text);
+  if (!jsonText) return recoverMalformedBareJsonToolCall(text);
+  try {
+    const parsed = JSON.parse(jsonText);
+    if (parsed.tool || parsed.name) {
+      const name = parsed.tool || parsed.name;
+      return normalizeGeneratedToolCall(name, parsed.input || parsed.arguments || parsed.parameters);
+    }
+    if (parsed.function?.name) {
+      return normalizeGeneratedToolCall(parsed.function.name, parsed.function.arguments || parsed.function.input || parsed.arguments);
+    }
+  } catch {
+    return recoverMalformedBareJsonToolCall(text);
+  }
+  return recoverMalformedBareJsonToolCall(text);
 }
 
 /**
@@ -52,50 +484,30 @@ function resolveModel(model) {
  * Uses Apple's MLX framework via Python subprocess for GPU-accelerated inference.
  */
 function createMLXProvider(config = {}) {
+  const configuredToolCallFormat = normalizeToolCallFormat(config.toolCallFormat);
   return {
     type: 'mlx',
 
     async chat({ model, messages, tools, system, maxTokens, temperature, signal } = {}) {
-      const mlxModel = resolveModel(model || config.model);
+      const modelSpec = resolveModelSpec(model || config.model, config);
+      const mlxModel = modelSpec?.repo;
       if (!mlxModel) throw new Error('MLX: no model specified');
-
-      // Build the prompt from messages
-      const chatMessages = [];
-      if (system) chatMessages.push({ role: 'system', content: system });
-      if (messages) {
-        for (const m of messages) {
-          if (typeof m.content === 'string') {
-            chatMessages.push({ role: m.role, content: m.content });
-          } else if (Array.isArray(m.content)) {
-            // Extract text from content blocks
-            const text = m.content
-              .filter(b => b.type === 'text')
-              .map(b => b.text)
-              .join('\n');
-            if (text) chatMessages.push({ role: m.role, content: text });
-            // Handle tool results
-            const toolResults = m.content.filter(b => b.type === 'tool_result');
-            for (const tr of toolResults) {
-              const resultText = typeof tr.content === 'string' ? tr.content
-                : Array.isArray(tr.content) ? tr.content.map(c => c.text || '').join('\n')
-                : JSON.stringify(tr.content);
-              chatMessages.push({ role: 'user', content: `[Tool result for ${tr.tool_use_id}]: ${resultText}` });
-            }
-          }
-        }
+      const adapterPath = modelSpec.adapterPath;
+      if (adapterPath && !fs.existsSync(adapterPath)) {
+        throw new Error(`MLX adapter path does not exist: ${adapterPath}`);
       }
 
-      // Build tool descriptions as system context (MLX doesn't natively support tool calling)
-      let toolContext = '';
-      if (tools && tools.length > 0) {
-        toolContext = '\n\nYou have access to these tools. To use a tool, respond with a JSON block:\n```json\n{"tool": "tool_name", "input": {…}}\n```\n\nAvailable tools:\n' +
-          tools.map(t => `- **${t.name}**: ${t.description}\n  Parameters: ${JSON.stringify(t.input_schema?.properties || {})}`).join('\n');
-        if (chatMessages.length > 0 && chatMessages[0].role === 'system') {
-          chatMessages[0].content += toolContext;
-        } else {
-          chatMessages.unshift({ role: 'system', content: toolContext.trim() });
-        }
-      }
+      const compactedChatMessages = prepareMessagesForMlxProvider({
+        messages,
+        tools: tools || [],
+        system,
+        compactDuplicates: true,
+        toolCallFormat: configuredToolCallFormat,
+      });
+      const toolContext = buildToolContext(tools || [], { toolCallFormat: configuredToolCallFormat });
+      const forcedGenerationPrefix = toolContext && configuredToolCallFormat === 'bare-json' && config.forceToolJsonPrefix
+        ? '{'
+        : '';
 
       // Python script for MLX inference
       const pyScript = `
@@ -105,12 +517,22 @@ from mlx_lm import load, generate
 messages = json.loads(sys.argv[1])
 max_tokens = int(sys.argv[2])
 model_name = sys.argv[3]
+adapter_path = sys.argv[4] if len(sys.argv) > 4 and sys.argv[4] else None
+forced_generation_prefix = sys.argv[5] if len(sys.argv) > 5 else ""
+
+load_kwargs = {}
+if adapter_path:
+    load_kwargs["adapter_path"] = adapter_path
 
-model, tokenizer = load(model_name)
+model, tokenizer = load(model_name, **load_kwargs)
 prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+if forced_generation_prefix:
+    prompt = prompt + forced_generation_prefix
 
 t0 = time.time()
 result = generate(model, tokenizer, prompt=prompt, max_tokens=max_tokens, verbose=False)
+if forced_generation_prefix:
+    result = forced_generation_prefix + result
 t1 = time.time()
 
 output_tokens = len(tokenizer.encode(result))
@@ -129,9 +551,9 @@ print(json.dumps({
 
       try {
         const { stdout, stderr } = await execAsync(
-          MLX_PYTHON,
-          ['-c', pyScript, JSON.stringify(chatMessages), String(maxTok), mlxModel],
-          { env: { ...process.env, TOKENIZERS_PARALLELISM: 'false' } }
+          resolveMlxPython(config),
+          ['-c', pyScript, JSON.stringify(compactedChatMessages), String(maxTok), mlxModel, adapterPath || '', forcedGenerationPrefix],
+          buildInferenceExecOptions(config, signal)
         );
 
         // Find JSON line in output
@@ -143,38 +565,39 @@ print(json.dumps({
         const latencyMs = Date.now() - start;
 
         // Parse tool calls from output if tools were provided
-        const content = [];
+        let content = data.text;
+        const toolCalls = [];
         let stopReason = 'end_turn';
 
-        if (tools && tools.length > 0) {
-          const toolMatch = data.text.match(/```json\s*\n?(\{[\s\S]*?\})\s*\n?```/);
-          if (toolMatch) {
-            try {
-              const toolCall = JSON.parse(toolMatch[1]);
-              if (toolCall.tool && toolCall.input) {
-                content.push({
-                  type: 'tool_use',
-                  id: `mlx_${Date.now()}`,
-                  name: toolCall.tool,
-                  input: toolCall.input,
-                });
-                stopReason = 'tool_use';
-              }
-            } catch { /* not a valid tool call */ }
+        if (toolContext) {
+          const toolCall = parseToolCallText(data.text);
+          if (toolCall?.name) {
+            toolCalls.push({
+              id: `mlx_${Date.now()}`,
+              name: toolCall.name,
+              input: toolCall.input,
+            });
+            content = '';
+            stopReason = 'tool_use';
           }
         }
 
-        if (content.length === 0) {
-          content.push({ type: 'text', text: data.text });
-        }
-
         return {
           content,
+          rawText: data.text,
+          toolCalls,
           stopReason,
           latencyMs,
+          usage: {
+            input: data.prompt_tokens,
+            output: data.output_tokens,
+            requestLatencyMs: latencyMs,
+          },
           inputTokens: data.prompt_tokens,
           outputTokens: data.output_tokens,
           model: mlxModel,
+          modelAlias: modelSpec.alias,
+          adapterPath,
           provider: 'mlx',
         };
       } catch (err) {
@@ -187,17 +610,23 @@ print(json.dumps({
       const models = [];
       const cacheDir = path.join(os.homedir(), '.cache', 'huggingface', 'hub');
 
-      for (const [alias, repo] of Object.entries(MLX_MODEL_REGISTRY)) {
+      for (const [alias, entry] of Object.entries(MLX_MODEL_REGISTRY)) {
+        const spec = normalizeRegistryEntry(entry);
+        if (!spec) continue;
+        const repo = spec.repo;
         // Check if model is cached locally
         const repoDir = `models--${repo.replace('/', '--')}`;
         const cached = fs.existsSync(path.join(cacheDir, repoDir));
         if (cached) {
+          const adapterInstalled = !spec.adapterPath || fs.existsSync(spec.adapterPath);
           models.push({
             id: alias,
             name: `${alias} (MLX)`,
             repo,
-            capabilities: ['code'],
-            cached: true,
+            adapterPath: spec.adapterPath || null,
+            adapterInstalled,
+            capabilities: spec.adapterPath ? ['code', 'tools'] : ['code'],
+            cached: adapterInstalled,
           });
         }
       }
@@ -206,9 +635,10 @@ print(json.dumps({
 
     async isAvailable() {
       // Check if MLX Python env exists
-      if (!fs.existsSync(MLX_PYTHON)) return false;
+      const mlxPython = resolveMlxPython(config);
+      if (!fs.existsSync(mlxPython)) return false;
       try {
-        const { stdout } = await execAsync(MLX_PYTHON, ['-c', 'import mlx_lm; print("ok")'], { timeout: 10000 });
+        const { stdout } = await execAsync(mlxPython, ['-c', 'import mlx_lm; print("ok")'], { timeout: 10000 });
         return stdout.trim() === 'ok';
       } catch {
         return false;
@@ -217,4 +647,37 @@ print(json.dumps({
   };
 }
 
-module.exports = { createMLXProvider, resolveModel, MLX_MODEL_REGISTRY, MLX_PYTHON, MLX_VENV };
+module.exports = {
+  createMLXProvider,
+  resolveModel,
+  resolveModelSpec,
+  resolveMlxPython,
+  resolveInferenceTimeoutMs,
+  buildInferenceExecOptions,
+  normalizeToolCallFormat,
+  normalizeToolSpec,
+  buildToolContext,
+  toolCallInput,
+  renderToolCallJsonBlock,
+  renderMlxProviderMessages,
+  prepareMessagesForMlxProvider,
+  compactDuplicateToolResults,
+  duplicateToolResultSignature,
+  extractJsonObject,
+  parseToolCallText,
+  MLX_MODEL_REGISTRY,
+  MLX_PYTHON,
+  MLX_VENV,
+  GEMMA_E4B_TOOLUSE_ADAPTER_PATH,
+  GEMMA_E4B_TOOLUSE_METADATA_PATH,
+  GEMMA_E4B_TOOLUSE_V2_ADAPTER_PATH,
+  GEMMA_E4B_TOOLUSE_V2_METADATA_PATH,
+  GEMMA_E4B_TOOLUSE_V3_ADAPTER_PATH,
+  GEMMA_E4B_TOOLUSE_V3_METADATA_PATH,
+  GEMMA_E4B_TOOLUSE_V4_ADAPTER_PATH,
+  GEMMA_E4B_TOOLUSE_V4_METADATA_PATH,
+  GEMMA_E4B_TOOLUSE_V5_ADAPTER_PATH,
+  GEMMA_E4B_TOOLUSE_V5_METADATA_PATH,
+  GEMMA_E4B_TOOLUSE_V6_ADAPTER_PATH,
+  GEMMA_E4B_TOOLUSE_V6_METADATA_PATH,
+};