modelmix 5.0.2 → 5.0.4

package/plugins/rlm/lib/planner-prompt.js

@@ -0,0 +1,137 @@
+const path = require('path');
+const { describeVariables } = require('./variable-descriptors');
+
+const PLANNER_SYSTEM_TEMPLATE = path.resolve(__dirname, '../prompts/planner.md');
+
+function positiveInteger(value, name) {
+    if (!Number.isInteger(value) || value <= 0) {
+        throw new TypeError(`${name} must be a positive integer.`);
+    }
+    return value;
+}
+
+function planningHint(name, descriptor, maxQueryBytes) {
+    const payloadBytes = descriptor.utf8Bytes ?? descriptor.estimatedBytes;
+    if (payloadBytes <= maxQueryBytes) {
+        return {
+            path: name,
+            strategy: 'direct',
+            reason: 'The complete serialized value fits within one query payload.'
+        };
+    }
+    if (descriptor.type === 'string') {
+        return {
+            path: name,
+            strategy: 'split-string-semantically',
+            boundary: 'paragraph',
+            reason: 'The string is larger than the maximum query payload.'
+        };
+    }
+    if (descriptor.type === 'array') {
+        const oversizedFields = [];
+        for (const [field, fieldDescriptor] of Object.entries(descriptor.itemShape.properties || {})) {
+            if (fieldDescriptor.stringSize?.utf8Bytes.max > maxQueryBytes) oversizedFields.push(field);
+        }
+        if (descriptor.itemSize.max > maxQueryBytes) {
+            return {
+                path: name,
+                strategy: 'split-oversized-items-semantically',
+                boundary: 'paragraph',
+                oversizedStringFields: oversizedFields,
+                reason: 'At least one array item is larger than the maximum query payload.'
+            };
+        }
+        const averageBytes = Math.max(1, descriptor.itemSize.average);
+        return {
+            path: name,
+            strategy: 'batch-array-items',
+            suggestedMaxItemsPerQuery: Math.max(1, Math.floor(maxQueryBytes / averageBytes)),
+            reason: 'The array is larger than one query payload while individual items fit.'
+        };
+    }
+    return {
+        path: name,
+        strategy: 'select-properties-or-descendants',
+        reason: 'The structured value is larger than the maximum query payload.'
+    };
+}
+
+function plannerTemplateData({
+    variables,
+    limits,
+    workerManifest,
+    outputMode = 'raw',
+    outputSchema = null
+}) {
+    if (!limits || typeof limits !== 'object' || Array.isArray(limits)) {
+        throw new TypeError('limits must be a plain object.');
+    }
+    const maxQueryBytes = positiveInteger(limits.maxQueryBytes, 'limits.maxQueryBytes');
+    const sandboxMemoryBytes = positiveInteger(limits.sandboxMemoryBytes, 'limits.sandboxMemoryBytes');
+    const maxConcurrentQueries = positiveInteger(
+        limits.maxConcurrentQueries,
+        'limits.maxConcurrentQueries'
+    );
+    const manifest = describeVariables(variables);
+    const planningHints = Object.entries(manifest.descriptors)
+        .map(([name, descriptor]) => planningHint(name, descriptor, maxQueryBytes));
+    if (!workerManifest || typeof workerManifest !== 'object' || Array.isArray(workerManifest)) {
+        throw new TypeError('workerManifest must be a plain object.');
+    }
+    if (!['message', 'json', 'block', 'raw'].includes(outputMode)) {
+        throw new TypeError(`Unsupported RLM planner output mode "${outputMode}".`);
+    }
+
+    return {
+        variableManifest: JSON.stringify(manifest, null, 2),
+        processingLimits: JSON.stringify({
+            maxQueryBytes,
+            sandboxMemoryBytes,
+            maxConcurrentQueries
+        }, null, 2),
+        planningHints: JSON.stringify(planningHints, null, 2),
+        workerManifest: JSON.stringify(workerManifest, null, 2),
+        outputRequirements: JSON.stringify({
+            mode: outputMode,
+            schema: outputSchema
+        }, null, 2),
+        maxQueryBytes,
+        maxConcurrentQueries
+    };
+}
+
+function createPlannerInvocation({
+    task,
+    variables,
+    limits,
+    workerManifest,
+    outputMode = 'raw',
+    outputSchema = null
+}) {
+    if (typeof task !== 'string' || task.trim().length === 0) {
+        throw new TypeError('task must be a non-empty string.');
+    }
+    return {
+        systemFile: PLANNER_SYSTEM_TEMPLATE,
+        assign: plannerTemplateData({
+            variables,
+            limits,
+            workerManifest,
+            outputMode,
+            outputSchema
+        }),
+        messages: [{
+            role: 'user',
+            content: [{ type: 'text', text: task }]
+        }],
+        plugins: { exclude: ['rlm'] },
+        history: false,
+        outputMode: 'raw'
+    };
+}
+
+module.exports = {
+    PLANNER_SYSTEM_TEMPLATE,
+    createPlannerInvocation,
+    plannerTemplateData
+};

package/plugins/rlm/lib/plugin.js

@@ -0,0 +1,203 @@
+const { RlmLimitError, validateRuntimeLimits } = require('./budget');
+const { createIsolatedVmSandbox } = require('./isolated-vm-sandbox');
+const { parseMarkdownDocument } = require('./markdown');
+const { createPlannerInvocation } = require('./planner-prompt');
+const {
+    RlmExecutionState,
+    createQueryRuntime,
+    requestTask,
+    sumTokens
+} = require('./runtime');
+const { createWorkerCatalog } = require('./worker-catalog');
+
+function isPlainObject(value) {
+    if (!value || typeof value !== 'object' || Array.isArray(value)) return false;
+    const prototype = Object.getPrototypeOf(value);
+    return prototype === Object.prototype || prototype === null;
+}
+
+function validateSandbox(sandbox) {
+    if (!sandbox || typeof sandbox !== 'object' || typeof sandbox.execute !== 'function') {
+        throw new TypeError('sandbox must define execute({ code, variables, query, limits }).');
+    }
+    return sandbox;
+}
+
+function validateDocuments(documents) {
+    if (documents === undefined) return {};
+    if (!isPlainObject(documents)) {
+        throw new TypeError('documents must be a plain object.');
+    }
+    for (const [name, document] of Object.entries(documents)) {
+        if (!isPlainObject(document)) {
+            throw new TypeError(`documents.${name} must be a plain object.`);
+        }
+        if (document.format !== 'markdown') {
+            throw new TypeError(`documents.${name}.format must be "markdown".`);
+        }
+        if (typeof document.content !== 'string') {
+            throw new TypeError(`documents.${name}.content must be a string.`);
+        }
+    }
+    return documents;
+}
+
+async function createExternalVariables(variables, documents) {
+    const result = { ...variables };
+    for (const [name, document] of Object.entries(documents)) {
+        if (Object.prototype.hasOwnProperty.call(result, name)) {
+            throw new TypeError(`External variable "${name}" is defined more than once.`);
+        }
+        result[name] = await parseMarkdownDocument(document.content);
+    }
+    return result;
+}
+
+function validateGeneratedCode(code) {
+    if (typeof code !== 'string' || code.trim().length === 0) {
+        throw new TypeError('RLM planner must return non-empty JavaScript code.');
+    }
+    if (/```|~~~/.test(code)) {
+        throw new SyntaxError('RLM planner returned Markdown fences instead of raw JavaScript.');
+    }
+    return code.trim();
+}
+
+function serializeResult(value) {
+    if (typeof value === 'string') return value;
+    let serialized;
+    try {
+        serialized = JSON.stringify(value);
+    } catch (error) {
+        throw new TypeError(`RLM sandbox result must be JSON-serializable: ${error.message}`);
+    }
+    if (serialized === undefined) {
+        throw new TypeError('RLM sandbox result must be JSON-serializable.');
+    }
+    return serialized;
+}
+
+function executionInput(context, configuredVariables) {
+    const invocation = context.request.config.rlmInvocation;
+    if (invocation !== undefined) {
+        if (!isPlainObject(invocation) || !(invocation.state instanceof RlmExecutionState)) {
+            throw new TypeError('Invalid internal RLM invocation state.');
+        }
+        return {
+            state: invocation.state,
+            task: invocation.task,
+            variables: invocation.variables
+        };
+    }
+    return {
+        state: null,
+        task: requestTask(context.request),
+        variables: configuredVariables
+    };
+}
+
+function runWithTimeout(operation, timeoutMs) {
+    let timeout;
+    return Promise.race([
+        operation(),
+        new Promise((_, reject) => {
+            timeout = setTimeout(() => reject(new RlmLimitError(
+                'maxWallTimeMs',
+                'RLM wall-time limit exceeded.'
+            )), timeoutMs);
+        })
+    ]).finally(() => clearTimeout(timeout));
+}
+
+function rlm({ maxDepth, variables = {}, documents, workers, limits, sandbox } = {}) {
+    if (!Number.isInteger(maxDepth) || maxDepth < 0) {
+        throw new TypeError('maxDepth must be a non-negative integer.');
+    }
+    if (!isPlainObject(variables)) {
+        throw new TypeError('variables must be a plain object.');
+    }
+    const validatedDocuments = validateDocuments(documents);
+    const catalog = createWorkerCatalog(workers);
+    const validatedLimits = validateRuntimeLimits(limits);
+    const sandboxAdapter = validateSandbox(
+        sandbox === undefined ? createIsolatedVmSandbox() : sandbox
+    );
+    const configuredVariables = createExternalVariables(variables, validatedDocuments);
+
+    return {
+        name: 'rlm',
+        async execute(context) {
+            if (context.request.outputMode === 'stream') {
+                throw new Error('RLM streaming is not supported; use a buffered output mode.');
+            }
+            const input = executionInput(context, await configuredVariables);
+            if (typeof input.task !== 'string' || input.task.trim().length === 0) {
+                throw new TypeError('RLM execution requires a non-empty task.');
+            }
+            if (!isPlainObject(input.variables)) {
+                throw new TypeError('RLM execution variables must be a plain object.');
+            }
+            const state = input.state || new RlmExecutionState(validatedLimits);
+            try {
+                const plannerStartedAt = Date.now();
+                const plannerResult = await state.budget.runPlanner(() => context.invoke(
+                    createPlannerInvocation({
+                        task: input.task,
+                        variables: input.variables,
+                        limits: validatedLimits,
+                        workerManifest: catalog.manifest,
+                        outputMode: context.request.outputMode,
+                        outputSchema: context.request.config.schema || null
+                    })
+                ));
+                state.record('planner', plannerResult, {
+                    worker: null,
+                    elapsedMs: Date.now() - plannerStartedAt
+                });
+                const code = validateGeneratedCode(plannerResult.message);
+                const query = createQueryRuntime({
+                    context,
+                    catalog,
+                    maxDepth,
+                    state
+                });
+                const timeoutMs = Math.max(
+                    1,
+                    validatedLimits.maxWallTimeMs - state.budget.snapshot().elapsedMs
+                );
+                const value = await runWithTimeout(
+                    () => sandboxAdapter.execute({
+                        code,
+                        variables: input.variables,
+                        query,
+                        limits: validatedLimits,
+                        execution: context.execution,
+                        timeoutMs
+                    }),
+                    timeoutMs
+                );
+                const message = serializeResult(value);
+                state.budget.accountFinalOutput(message);
+                return {
+                    message,
+                    tokens: sumTokens(state.calls),
+                    rlm: state.diagnostics(context.execution)
+                };
+            } catch (error) {
+                error.rlm = state.diagnostics(
+                    context.execution,
+                    error.limit ? `limit:${error.limit}` : 'error'
+                );
+                throw error;
+            }
+        }
+    };
+}
+
+module.exports = {
+    createExternalVariables,
+    rlm,
+    serializeResult,
+    validateDocuments,
+    validateGeneratedCode
+};

package/plugins/rlm/lib/runtime.js

@@ -0,0 +1,146 @@
+const { createRuntimeBudget } = require('./budget');
+
+function textContent(content) {
+    if (typeof content === 'string') return content;
+    if (!Array.isArray(content)) return '';
+    return content
+        .filter(part => part?.type === 'text' && typeof part.text === 'string')
+        .map(part => part.text)
+        .join('\n');
+}
+
+function requestTask(request) {
+    return request.messages
+        .filter(message => message?.role === 'user')
+        .map(message => textContent(message.content))
+        .filter(Boolean)
+        .join('\n\n');
+}
+
+function queryPayloadBytes({ system, message }) {
+    return Buffer.byteLength(JSON.stringify({ system, message }), 'utf8');
+}
+
+function validateQueryInput(input) {
+    if (!input || typeof input !== 'object' || Array.isArray(input)) {
+        throw new TypeError('query() expects a plain object.');
+    }
+    if (typeof input.worker !== 'string' || input.worker.length === 0) {
+        throw new TypeError('query.worker must be a non-empty string.');
+    }
+    if (typeof input.system !== 'string' || input.system.trim().length === 0) {
+        throw new TypeError('query.system must be a non-empty string.');
+    }
+    if (typeof input.message !== 'string') {
+        throw new TypeError('query.message must be a string.');
+    }
+    return input;
+}
+
+function sumTokens(calls) {
+    const totals = {
+        input: 0,
+        output: 0,
+        total: 0,
+        cached: 0,
+        cacheWrite: 0,
+        cacheWrite5m: 0,
+        cacheWrite1h: 0,
+        uncachedInput: 0,
+        cost: 0
+    };
+    let found = false;
+    for (const call of calls) {
+        const tokens = call.tokens;
+        if (!tokens) continue;
+        found = true;
+        for (const key of Object.keys(totals)) {
+            if (Number.isFinite(tokens[key])) totals[key] += tokens[key];
+        }
+    }
+    return found ? totals : undefined;
+}
+
+class RlmExecutionState {
+    constructor(limits) {
+        this.budget = createRuntimeBudget(limits);
+        this.calls = [];
+    }
+
+    record(kind, result, details = {}, { includeTokens = true } = {}) {
+        this.calls.push({
+            kind,
+            ...details,
+            execution: result.execution || null,
+            tokens: includeTokens ? (result.tokens || null) : null
+        });
+    }
+
+    diagnostics(rootExecution, terminationReason = 'completed') {
+        return {
+            execution: rootExecution,
+            calls: this.calls.map(call => ({ ...call })),
+            budget: this.budget.snapshot(),
+            terminationReason
+        };
+    }
+}
+
+function createQueryRuntime({ context, catalog, maxDepth, state }) {
+    return async rawInput => {
+        const startedAt = Date.now();
+        const input = validateQueryInput(rawInput);
+        const model = catalog.get(input.worker);
+        const directLeaf = context.execution.depth >= maxDepth;
+        const payloadBytes = queryPayloadBytes(input);
+        state.budget.assertQueryPayload(payloadBytes);
+        const invocation = directLeaf
+            ? {
+                model,
+                system: input.system,
+                messages: [{ role: 'user', content: input.message }],
+                plugins: { exclude: ['rlm'] },
+                history: false,
+                outputMode: 'raw'
+            }
+            : {
+                model,
+                messages: [{ role: 'user', content: input.system }],
+                config: {
+                    rlmInvocation: {
+                        state,
+                        task: input.system,
+                        variables: { input: input.message }
+                    },
+                },
+                plugins: 'inherit',
+                history: false,
+                outputMode: 'raw'
+            };
+        const result = directLeaf
+            ? await state.budget.runQuery(
+                { payloadBytes },
+                () => context.invoke(invocation)
+            )
+            : await context.invoke(invocation);
+        if (typeof result.message !== 'string') {
+            throw new TypeError(`RLM worker "${input.worker}" must return a message string.`);
+        }
+        state.record('worker', result, {
+            worker: input.worker,
+            directLeaf,
+            depthBoundary: directLeaf ? maxDepth : null,
+            elapsedMs: Date.now() - startedAt
+        }, { includeTokens: directLeaf });
+        return result.message;
+    };
+}
+
+module.exports = {
+    RlmExecutionState,
+    createQueryRuntime,
+    queryPayloadBytes,
+    requestTask,
+    sumTokens,
+    validateQueryInput
+};