agentfootprint/dist/memory/stages/pickByBudget.js

The explainable agent framework — backtrack a wrong answer to the exact context that caused it (evidence, not guesses). Built on footprintjs.

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.pickByBudget = void 0;
const footprintjs_1 = require("footprintjs");
const tokenize_js_1 = require("./tokenize.js");
const provenance_js_1 = require("../retrieval/provenance.js");
const DEFAULT_RESERVE = 256;
const DEFAULT_MINIMUM = 100;
/**
 * Build the decider function. Returns the full `DecisionResult` so
 * DeciderHandler recognizes the DECISION_RESULT brand and attaches
 * evidence to `FlowRecorder.onDecision`. Predicates read from `scope`
 * (not closed-over locals) so the temp recorder captures the values
 * that drove the choice.
 */
function buildPickDecider(config) {
    const reserveTokens = config.reserveTokens ?? DEFAULT_RESERVE;
    const minimumTokens = config.minimumTokens ?? DEFAULT_MINIMUM;
    return (scope) => (0, footprintjs_1.decide)(scope, [
        {
            when: (s) => (s.loaded ?? []).length === 0,
            then: 'skip-empty',
            label: 'no entries loaded — nothing to pick',
        },
        {
            when: (s) => (s.contextTokensRemaining ?? 0) - reserveTokens < minimumTokens,
            then: 'skip-no-budget',
            label: 'budget below minimum threshold — skip injection',
        },
    ], 'pick');
}
/**
 * Both skip branches share the same body — no entries survive.
 *
 * When a retrieval record is present, every candidate that had been
 * admitted by the retrieval strategy is re-marked as rejected here, with
 * the budget as the reason. A record that still claimed those chunks were
 * admitted, after a branch that injected nothing, would be a record that
 * disagrees with the prompt.
 */
const skipStage = (scope) => {
    scope.selected = [];
    reviseAdmitted(scope, () => false, undefined, 0);
};
/**
 * Re-mark the retrieval record against what the budget picker actually
 * kept. `keptIds` is consulted only for candidates the strategy already
 * admitted — a below-threshold candidate keeps its own reason, because
 * the budget never got the chance to reject it.
 *
 * `charsKept` re-states the retriever's `maxChars` spend against what
 * survived here (8.19.0). The retrieval recorded what IT spent; if this
 * picker then drops an entry for tokens, a `charsUsed` still counting that
 * entry's passage would be a record disagreeing with its own
 * `admittedCount`. It is written only when the retriever set a budget —
 * nothing new appears on a record that never had one.
 */
function reviseAdmitted(scope, wasKept, overCap, charsKept) {
    const evidence = scope.retrieved;
    if (!evidence?.candidates)
        return;
    let admittedCount = 0;
    const candidates = evidence.candidates.map((c) => {
        if (!c.admitted)
            return c;
        if (wasKept(c.id)) {
            admittedCount += 1;
            return c;
        }
        const reason = overCap?.(c.id) ? 'over-max-entries' : 'over-budget';
        return { ...c, admitted: false, reason };
    });
    scope.retrieved = {
        ...evidence,
        candidates,
        admittedCount,
        rejectedCount: candidates.length - admittedCount,
        ...(evidence.maxChars !== undefined && { charsUsed: charsKept }),
    };
}
/** The `pick` branch: greedy newest-first selection within budget. */
function buildPickStage(config) {
    const reserveTokens = config.reserveTokens ?? DEFAULT_RESERVE;
    const countTokens = config.countTokens ?? tokenize_js_1.approximateTokenCounter;
    const maxEntries = config.maxEntries;
    return (scope) => {
        const loaded = scope.loaded ?? [];
        const budget = (scope.contextTokensRemaining ?? 0) - reserveTokens;
        // Sort newest-first. Secondary key on `id` guarantees deterministic
        // ordering when entries share `updatedAt` (batch writes, low-resolution
        // clocks) — without it, ties resolve to implementation-defined order
        // which breaks trace replay and A/B eval comparisons.
        const byNewest = [...loaded].sort((a, b) => {
            const byTime = b.updatedAt - a.updatedAt;
            if (byTime !== 0)
                return byTime;
            return a.id < b.id ? -1 : a.id > b.id ? 1 : 0;
        });
        const picked = [];
        const overCap = new Set();
        let used = 0;
        for (const entry of byNewest) {
            if (maxEntries !== undefined && picked.length >= maxEntries) {
                // Everything from here on is refused by the cap, not by the
                // budget. The two reasons are not interchangeable: one is fixed
                // by shortening the corpus, the other by raising `maxEntries`.
                //
                // This was a `break` before 8.8.0. It is a `continue` so the
                // remaining entries can be NAMED as cap-rejected — and it cannot
                // change what is picked, because once the cap is reached the
                // condition stays true for every later iteration.
                overCap.add(entry.id);
                continue;
            }
            const cost = (0, tokenize_js_1.countMessageTokens)(entry.value, countTokens);
            if (used + cost > budget)
                continue; // skip this entry, try smaller ones
            picked.push(entry);
            used += cost;
        }
        // Emit in chronological order — `picked` is newest-first; reverse.
        scope.selected = picked.reverse();
        const kept = new Set(picked.map((e) => e.id));
        // Measured through the SAME accessor the retriever spent its character
        // budget with, so the two numbers are the same kind of number.
        const charsKept = picked.reduce((total, entry) => total + (0, provenance_js_1.chunkText)(entry.value).length, 0);
        reviseAdmitted(scope, (id) => kept.has(id), (id) => overCap.has(id), charsKept);
    };
}
/**
 * Append the pick-by-budget decider + branches to `builder`. Returns
 * the builder so calls chain naturally:
 *
 * ```ts
 * let b = flowChart<MemoryState>('LoadRecent', loadRecent(config), 'load-recent');
 * b = pickByBudget(pickConfig)(b);
 * b = b.addFunction('Format', formatDefault(formatConfig), 'format-default');
 * ```
 *
 * Generic in `T` so consumers whose scope extends `MemoryState` (e.g.,
 * an AgentLoopState that embeds memory fields) can compose this into
 * their own pipeline without casting.
 */
function pickByBudget(config = {}) {
    const decider = buildPickDecider(config);
    const pickStage = buildPickStage(config);
    return (builder) => {
        return builder
            .addDeciderFunction('PickDecider', decider, 'pick-decider', 'Decide whether to pick entries, skip (empty), or skip (no budget)')
            .addFunctionBranch('skip-empty', 'SkipEmpty', skipStage, 'Mark selected as [] — no entries loaded')
            .addFunctionBranch('skip-no-budget', 'SkipNoBudget', skipStage, 'Mark selected as [] — budget below minimum')
            .addFunctionBranch('pick', 'Pick', pickStage, 'Greedy newest-first selection within token budget')
            .end();
    };
}
exports.pickByBudget = pickByBudget;
//# sourceMappingURL=pickByBudget.js.map