@ax-llm/ax 9.0.21 → 9.0.22

package/cjs/dsp/generate.js

@@ -0,0 +1,314 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.AxGenerate = void 0;
+const web_1 = require("stream/web");
+const util_js_1 = require("../ai/util.js");
+const functions_js_1 = require("../funcs/functions.js");
+const index_js_1 = require("../mem/index.js");
+const index_js_2 = require("../trace/index.js");
+const asserts_js_1 = require("./asserts.js");
+const extract_js_1 = require("./extract.js");
+const program_js_1 = require("./program.js");
+const prompt_js_1 = require("./prompt.js");
+class AxGenerate extends program_js_1.AxProgram {
+    pt;
+    asserts;
+    streamingAsserts;
+    options;
+    funcProc;
+    functionList;
+    constructor(ai, signature, options) {
+        super(ai, signature);
+        this.options = options;
+        this.pt = new (options?.promptTemplate ?? prompt_js_1.AxPromptTemplate)(this.signature);
+        this.asserts = this.options?.asserts ?? [];
+        this.streamingAsserts = this.options?.streamingAsserts ?? [];
+        this.functionList = this.options?.functions?.map((f) => f.name).join(', ');
+        this.usage = [];
+        if (this.options?.functions) {
+            this.funcProc = new functions_js_1.AxFunctionProcessor(this.options?.functions);
+            this.updateSigForFunctions();
+        }
+    }
+    updateSigForFunctions = () => {
+        // AI supports function calling natively so
+        // no need to add fields for function call
+        if (this.ai.getFeatures().functions) {
+            return;
+        }
+        // These are the fields for the function call only needed when the underlying LLM API does not support function calling natively in the API.
+        this.signature.addOutputField({
+            name: 'functionName',
+            description: 'Name of function to call',
+            isOptional: true
+        });
+        this.signature.addOutputField({
+            name: 'functionArguments',
+            description: 'Arguments of function to call',
+            isOptional: true
+        });
+    };
+    addAssert = (fn, message, optional) => {
+        this.asserts.push({ fn, message, optional });
+    };
+    addStreamingAssert = (fieldName, fn, message, optional) => {
+        this.streamingAsserts.push({ fieldName, fn, message, optional });
+    };
+    async forwardSendRequest({ mem, sessionId, traceId, ai, modelConfig: mc, stream, model, rateLimiter }) {
+        const chatPrompt = mem?.history(sessionId) ?? [];
+        if (chatPrompt.length === 0) {
+            throw new Error('No chat prompt found');
+        }
+        const functions = this.options?.functions;
+        const functionCall = this.options?.functionCall;
+        const hasJSON = this.signature
+            .getOutputFields()
+            .some((f) => f?.type?.name === 'json' || f?.type?.isArray);
+        const modelConfig = mc
+            ? {
+                ...mc,
+                ...(hasJSON ? { outputFormat: 'json_object' } : {})
+            }
+            : undefined;
+        const res = await ai.chat({
+            chatPrompt,
+            functions,
+            functionCall,
+            modelConfig,
+            ...(model ? { model } : {})
+        }, {
+            ...(sessionId ? { sessionId } : {}),
+            ...(traceId ? { traceId } : {}),
+            ...(rateLimiter ? { rateLimiter } : {}),
+            stream
+        });
+        return res;
+    }
+    async forwardCore({ mem, sessionId, traceId, ai, modelConfig, rateLimiter, stream = false }) {
+        const usageInfo = {
+            ai: this.ai.getName(),
+            model: this.ai.getModelInfo().name
+        };
+        const res = await this.forwardSendRequest({
+            mem,
+            sessionId,
+            traceId,
+            ai,
+            stream,
+            modelConfig,
+            rateLimiter
+        });
+        if (res instanceof web_1.ReadableStream) {
+            return (await this.processSteamingResponse({
+                res,
+                usageInfo,
+                mem,
+                traceId,
+                sessionId
+            }));
+        }
+        return (await this.processResponse({
+            res,
+            usageInfo,
+            mem,
+            traceId,
+            sessionId
+        }));
+    }
+    async processSteamingResponse({ res, usageInfo, mem, sessionId, traceId }) {
+        const functionCalls = [];
+        const values = {};
+        const xstate = { s: -1 };
+        let content = '';
+        for await (const v of res) {
+            for (const result of v.results ?? []) {
+                if (v.modelUsage) {
+                    this.usage.push({ ...usageInfo, ...v.modelUsage });
+                }
+                if (result.content) {
+                    content += result.content;
+                    mem.updateResult({ name: result.name, content }, sessionId);
+                    (0, asserts_js_1.assertStreamingAssertions)(this.streamingAsserts, values, xstate, content);
+                    (0, extract_js_1.streamingExtractValues)(this.signature, values, xstate, content);
+                    (0, asserts_js_1.assertAssertions)(this.asserts, values);
+                }
+                if (result.functionCalls) {
+                    (0, util_js_1.mergeFunctionCalls)(functionCalls, result.functionCalls);
+                    mem.updateResult({ name: result.name, content, functionCalls }, sessionId);
+                }
+                if (result.finishReason === 'length') {
+                    throw new Error('Max tokens reached before completion');
+                }
+            }
+        }
+        const funcs = parseFunctions(this.ai, functionCalls, values);
+        if (funcs) {
+            await this.processFunctions(funcs, mem, sessionId, traceId);
+        }
+        (0, extract_js_1.streamingExtractFinalValue)(values, xstate, content);
+        (0, asserts_js_1.assertAssertions)(this.asserts, values);
+        return { ...values };
+    }
+    async processResponse({ res, usageInfo, mem, sessionId, traceId }) {
+        const values = {};
+        for (const result of res.results ?? []) {
+            if (res.modelUsage) {
+                this.usage.push({ ...usageInfo, ...res.modelUsage });
+            }
+            mem.addResult(result, sessionId);
+            if (result.content) {
+                (0, extract_js_1.extractValues)(this.signature, values, result.content);
+                (0, asserts_js_1.assertAssertions)(this.asserts, values);
+            }
+            if (result.functionCalls) {
+                const funcs = parseFunctions(this.ai, result.functionCalls, values);
+                if (funcs) {
+                    await this.processFunctions(funcs, mem, sessionId, traceId);
+                }
+            }
+            if (result.finishReason === 'length') {
+                throw new Error('Max tokens reached before completion');
+            }
+        }
+        return { ...values };
+    }
+    async _forward(values, options, span) {
+        const maxRetries = this.options?.maxRetries ?? options?.maxRetries ?? 5;
+        const maxSteps = this.options?.maxSteps ?? options?.maxSteps ?? 10;
+        const mem = this.options?.mem ?? options?.mem ?? new index_js_1.AxMemory();
+        const canStream = this.ai.getFeatures().streaming;
+        let err;
+        if (this.sigHash !== this.signature.hash()) {
+            const promptTemplate = this.options?.promptTemplate ?? prompt_js_1.AxPromptTemplate;
+            this.pt = new promptTemplate(this.signature);
+        }
+        const prompt = this.pt.render(values, {
+            examples: this.examples,
+            demos: this.demos
+        });
+        const userMsg = { role: 'user', content: prompt };
+        mem.add(userMsg, options?.sessionId);
+        multiStepLoop: for (let n = 0; n < maxSteps; n++) {
+            for (let i = 0; i < maxRetries; i++) {
+                try {
+                    const { sessionId, traceId, modelConfig, stream: doStream } = options ?? {};
+                    const stream = canStream && doStream;
+                    const output = await this.forwardCore({
+                        ai: options?.ai ?? this.ai,
+                        mem,
+                        sessionId,
+                        traceId,
+                        modelConfig,
+                        stream,
+                        maxSteps: options?.maxSteps,
+                        rateLimiter: options?.rateLimiter
+                    });
+                    const lastMemItem = mem.getLast(sessionId);
+                    if (lastMemItem?.role === 'function') {
+                        continue multiStepLoop;
+                    }
+                    (0, asserts_js_1.assertRequiredFields)(this.signature, output);
+                    this.trace = { ...output };
+                    return output;
+                }
+                catch (e) {
+                    let extraFields;
+                    span?.recordAxSpanException(e);
+                    if (e instanceof extract_js_1.ValidationError) {
+                        extraFields = e.getFixingInstructions();
+                        err = e;
+                    }
+                    else if (e instanceof asserts_js_1.AxAssertionError) {
+                        const e1 = e;
+                        extraFields = e1.getFixingInstructions(this.signature);
+                        err = e;
+                    }
+                    else {
+                        throw e;
+                    }
+                    if (extraFields) {
+                        const content = this.pt.renderExtraFields(extraFields);
+                        mem.add({ role: 'user', content }, options?.sessionId);
+                        if (options?.debug) {
+                            console.log('Error Correction:', content);
+                        }
+                    }
+                }
+            }
+            if (err instanceof asserts_js_1.AxAssertionError && err.getOptional()) {
+                return err.getValue();
+            }
+            throw new Error(`Unable to fix validation error: ${err?.message}`);
+        }
+        throw new Error('Could not complete task within maximum allowed steps');
+    }
+    async forward(values, options) {
+        const tracer = this.options?.tracer ?? options?.tracer;
+        if (!tracer) {
+            return await this._forward(values, options);
+        }
+        const attributes = {
+            ['generate.signature']: this.signature.toString(),
+            ['generate.functions']: this.functionList ?? 'none'
+        };
+        return await tracer.startActiveSpan('Generate', {
+            kind: index_js_2.AxSpanKind.SERVER,
+            attributes
+        }, async (span) => {
+            const res = this._forward(values, options, span);
+            span.end();
+            return res;
+        });
+    }
+    processFunctions = async (functionCalls, mem, sessionId, traceId) => {
+        // Map each function call to a promise that resolves to the function result or null
+        const promises = functionCalls.map((func) => this.funcProc?.execute(func, { sessionId, traceId }).then((fres) => {
+            if (fres?.id) {
+                return {
+                    role: 'function',
+                    result: fres.result ?? '',
+                    functionId: fres.id
+                };
+            }
+            return null; // Returning null for function calls that don't meet the condition
+        }));
+        // Wait for all promises to resolve
+        const results = await Promise.all(promises);
+        results.forEach((result) => {
+            if (result) {
+                mem.add(result, sessionId);
+            }
+        });
+    };
+}
+exports.AxGenerate = AxGenerate;
+function parseFunctions(ai, functionCalls, values) {
+    if (!functionCalls || functionCalls.length === 0) {
+        return;
+    }
+    if (ai.getFeatures().functions) {
+        const funcs = functionCalls.map((f) => ({
+            id: f.id,
+            name: f.function.name,
+            args: f.function.params
+        }));
+        // for (const [i, f] of funcs.entries()) {
+        //   values['functionName' + i] = f.name;
+        //   values['functionArguments' + i] =
+        //     typeof f.args === 'object' ? JSON.stringify(f.args) : f.args;
+        // }
+        return funcs;
+    }
+    else if (values.functionName) {
+        const { functionName, functionArguments } = values;
+        delete values.functionName;
+        delete values.functionArguments;
+        return [
+            {
+                name: functionName,
+                args: functionArguments
+            }
+        ];
+    }
+}
+//# sourceMappingURL=generate.js.map

package/cjs/dsp/generate.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"generate.js","sourceRoot":"","sources":["../../../dsp/generate.ts"],"names":[],"mappings":";;;AAAA,oCAA4C;AAU5C,2CAAmD;AACnD,wDAG+B;AAC/B,8CAA4D;AAC5D,gDAA2E;AAE3E,6CAOsB;AACtB,6CAMsB;AACtB,6CAKsB;AACtB,2CAA+C;AAgC/C,MAAa,UAGX,SAAQ,sBAAkB;IAClB,EAAE,CAAmB;IACrB,OAAO,CAAgB;IACvB,gBAAgB,CAAyB;IACzC,OAAO,CAAqB;IAC5B,QAAQ,CAAuB;IAC/B,YAAY,CAAU;IAE9B,YACE,EAAe,EACf,SAAyC,EACzC,OAAqC;QAErC,KAAK,CAAC,EAAE,EAAE,SAAS,CAAC,CAAC;QAErB,IAAI,CAAC,OAAO,GAAG,OAAO,CAAC;QACvB,IAAI,CAAC,EAAE,GAAG,IAAI,CAAC,OAAO,EAAE,cAAc,IAAI,4BAAgB,CAAC,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC;QAC5E,IAAI,CAAC,OAAO,GAAG,IAAI,CAAC,OAAO,EAAE,OAAO,IAAI,EAAE,CAAC;QAC3C,IAAI,CAAC,gBAAgB,GAAG,IAAI,CAAC,OAAO,EAAE,gBAAgB,IAAI,EAAE,CAAC;QAC7D,IAAI,CAAC,YAAY,GAAG,IAAI,CAAC,OAAO,EAAE,SAAS,EAAE,GAAG,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,IAAI,CAAC,CAAC,IAAI,CAAC,IAAI,CAAC,CAAC;QAC3E,IAAI,CAAC,KAAK,GAAG,EAAE,CAAC;QAEhB,IAAI,IAAI,CAAC,OAAO,EAAE,SAAS,EAAE,CAAC;YAC5B,IAAI,CAAC,QAAQ,GAAG,IAAI,kCAAmB,CAAC,IAAI,CAAC,OAAO,EAAE,SAAS,CAAC,CAAC;YACjE,IAAI,CAAC,qBAAqB,EAAE,CAAC;QAC/B,CAAC;IACH,CAAC;IAEO,qBAAqB,GAAG,GAAG,EAAE;QACnC,2CAA2C;QAC3C,0CAA0C;QAC1C,IAAI,IAAI,CAAC,EAAE,CAAC,WAAW,EAAE,CAAC,SAAS,EAAE,CAAC;YACpC,OAAO;QACT,CAAC;QAED,4IAA4I;QAC5I,IAAI,CAAC,SAAS,CAAC,cAAc,CAAC;YAC5B,IAAI,EAAE,cAAc;YACpB,WAAW,EAAE,0BAA0B;YACvC,UAAU,EAAE,IAAI;SACjB,CAAC,CAAC;QAEH,IAAI,CAAC,SAAS,CAAC,cAAc,CAAC;YAC5B,IAAI,EAAE,mBAAmB;YACzB,WAAW,EAAE,+BAA+B;YAC5C,UAAU,EAAE,IAAI;SACjB,CAAC,CAAC;IACL,CAAC,CAAC;IAEK,SAAS,GAAG,CACjB,EAAqB,EACrB,OAAgB,EAChB,QAAkB,EAClB,EAAE;QACF,IAAI,CAAC,OAAO,CAAC,IAAI,CAAC,EAAE,EAAE,EAAE,OAAO,EAAE,QAAQ,EAAE,CAAC,CAAC;IAC/C,CAAC,CAAC;IAEK,kBAAkB,GAAG,CAC1B,SAAiB,EACjB,EAA8B,EAC9B,OAAgB,EAChB,QAAkB,EAClB,EAAE;QACF,IAAI,CAAC,gBAAgB,CAAC,IAAI,CAAC,EAAE,SAAS,EAAE,EAAE,EAAE,OAAO,EAAE,QAAQ,EAAE,CAAC,CAAC;IACnE,CAAC,CAAC;IAEM,KAAK,CAAC,kBAAkB,CAAC,EAC/B,GAAG,EACH,SAAS,EACT,OAAO,EACP,EAAE,EACF,WAAW,EAAE,EAAE,EACf,MAAM,EACN,KAAK,EACL,WAAW,EAGZ;QACC,MAAM,UAAU,GAAG,GAAG,EAAE,OAAO,CAAC,SAAS,CAAC,IAAI,EAAE,CAAC;QAEjD,IAAI,UAAU,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;YAC5B,MAAM,IAAI,KAAK,CAAC,sBAAsB,CAAC,CAAC;QAC1C,CAAC;QAED,MAAM,SAAS,GAAG,IAAI,CAAC,OAAO,EAAE,SAAS,CAAC;QAC1C,MAAM,YAAY,GAAG,IAAI,CAAC,OAAO,EAAE,YAAY,CAAC;QAEhD,MAAM,OAAO,GAAG,IAAI,CAAC,SAAS;aAC3B,eAAe,EAAE;aACjB,IAAI,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC,EAAE,IAAI,EAAE,IAAI,KAAK,MAAM,IAAI,CAAC,EAAE,IAAI,EAAE,OAAO,CAAC,CAAC;QAE7D,MAAM,WAAW,GAAG,EAAE;YACpB,CAAC,CAAC;gBACE,GAAG,EAAE;gBACL,GAAG,CAAC,OAAO,CAAC,CAAC,CAAC,EAAE,YAAY,EAAE,aAAa,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;aACpD;YACH,CAAC,CAAC,SAAS,CAAC;QAEd,MAAM,GAAG,GAAG,MAAM,EAAE,CAAC,IAAI,CACvB;YACE,UAAU;YACV,SAAS;YACT,YAAY;YACZ,WAAW;YACX,GAAG,CAAC,KAAK,CAAC,CAAC,CAAC,EAAE,KAAK,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;SAC5B,EACD;YACE,GAAG,CAAC,SAAS,CAAC,CAAC,CAAC,EAAE,SAAS,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;YACnC,GAAG,CAAC,OAAO,CAAC,CAAC,CAAC,EAAE,OAAO,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;YAC/B,GAAG,CAAC,WAAW,CAAC,CAAC,CAAC,EAAE,WAAW,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC;YACvC,MAAM;SACP,CACF,CAAC;QAEF,OAAO,GAAG,CAAC;IACb,CAAC;IAEO,KAAK,CAAC,WAAW,CAAC,EACxB,GAAG,EACH,SAAS,EACT,OAAO,EACP,EAAE,EACF,WAAW,EACX,WAAW,EACX,MAAM,GAAG,KAAK,EAMf;QACC,MAAM,SAAS,GAAG;YAChB,EAAE,EAAE,IAAI,CAAC,EAAE,CAAC,OAAO,EAAE;YACrB,KAAK,EAAE,IAAI,CAAC,EAAE,CAAC,YAAY,EAAE,CAAC,IAAI;SACnC,CAAC;QAEF,MAAM,GAAG,GAAG,MAAM,IAAI,CAAC,kBAAkB,CAAC;YACxC,GAAG;YACH,SAAS;YACT,OAAO;YACP,EAAE;YACF,MAAM;YACN,WAAW;YACX,WAAW;SACZ,CAAC,CAAC;QAEH,IAAI,GAAG,YAAY,oBAAc,EAAE,CAAC;YAClC,OAAO,CAAC,MAAM,IAAI,CAAC,uBAAuB,CAAC;gBACzC,GAAG;gBACH,SAAS;gBACT,GAAG;gBACH,OAAO;gBACP,SAAS;aACV,CAAC,CAAmB,CAAC;QACxB,CAAC;QAED,OAAO,CAAC,MAAM,IAAI,CAAC,eAAe,CAAC;YACjC,GAAG;YACH,SAAS;YACT,GAAG;YACH,OAAO;YACP,SAAS;SACV,CAAC,CAAmB,CAAC;IACxB,CAAC;IAEO,KAAK,CAAC,uBAAuB,CAAC,EACpC,GAAG,EACH,SAAS,EACT,GAAG,EACH,SAAS,EACT,OAAO,EAGR;QACC,MAAM,aAAa,GACjB,EAAE,CAAC;QACL,MAAM,MAAM,GAAG,EAAE,CAAC;QAClB,MAAM,MAAM,GAAoB,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC;QAE1C,IAAI,OAAO,GAAG,EAAE,CAAC;QAEjB,IAAI,KAAK,EAAE,MAAM,CAAC,IAAI,GAAG,EAAE,CAAC;YAC1B,KAAK,MAAM,MAAM,IAAI,CAAC,CAAC,OAAO,IAAI,EAAE,EAAE,CAAC;gBACrC,IAAI,CAAC,CAAC,UAAU,EAAE,CAAC;oBACjB,IAAI,CAAC,KAAK,CAAC,IAAI,CAAC,EAAE,GAAG,SAAS,EAAE,GAAG,CAAC,CAAC,UAAU,EAAE,CAAC,CAAC;gBACrD,CAAC;gBAED,IAAI,MAAM,CAAC,OAAO,EAAE,CAAC;oBACnB,OAAO,IAAI,MAAM,CAAC,OAAO,CAAC;oBAE1B,GAAG,CAAC,YAAY,CAAC,EAAE,IAAI,EAAE,MAAM,CAAC,IAAI,EAAE,OAAO,EAAE,EAAE,SAAS,CAAC,CAAC;oBAE5D,IAAA,sCAAyB,EACvB,IAAI,CAAC,gBAAgB,EACrB,MAAM,EACN,MAAM,EACN,OAAO,CACR,CAAC;oBACF,IAAA,mCAAsB,EAAC,IAAI,CAAC,SAAS,EAAE,MAAM,EAAE,MAAM,EAAE,OAAO,CAAC,CAAC;oBAChE,IAAA,6BAAgB,EAAC,IAAI,CAAC,OAAO,EAAE,MAAM,CAAC,CAAC;gBACzC,CAAC;gBAED,IAAI,MAAM,CAAC,aAAa,EAAE,CAAC;oBACzB,IAAA,4BAAkB,EAAC,aAAa,EAAE,MAAM,CAAC,aAAa,CAAC,CAAC;oBAExD,GAAG,CAAC,YAAY,CACd,EAAE,IAAI,EAAE,MAAM,CAAC,IAAI,EAAE,OAAO,EAAE,aAAa,EAAE,EAC7C,SAAS,CACV,CAAC;gBACJ,CAAC;gBAED,IAAI,MAAM,CAAC,YAAY,KAAK,QAAQ,EAAE,CAAC;oBACrC,MAAM,IAAI,KAAK,CAAC,sCAAsC,CAAC,CAAC;gBAC1D,CAAC;YACH,CAAC;QACH,CAAC;QAED,MAAM,KAAK,GAAG,cAAc,CAAC,IAAI,CAAC,EAAE,EAAE,aAAa,EAAE,MAAM,CAAC,CAAC;QAC7D,IAAI,KAAK,EAAE,CAAC;YACV,MAAM,IAAI,CAAC,gBAAgB,CAAC,KAAK,EAAE,GAAG,EAAE,SAAS,EAAE,OAAO,CAAC,CAAC;QAC9D,CAAC;QAED,IAAA,uCAA0B,EAAC,MAAM,EAAE,MAAM,EAAE,OAAO,CAAC,CAAC;QACpD,IAAA,6BAAgB,EAAC,IAAI,CAAC,OAAO,EAAE,MAAM,CAAC,CAAC;QAEvC,OAAO,EAAE,GAAG,MAAM,EAAoB,CAAC;IACzC,CAAC;IAEO,KAAK,CAAC,eAAe,CAAC,EAC5B,GAAG,EACH,SAAS,EACT,GAAG,EACH,SAAS,EACT,OAAO,EACyC;QAChD,MAAM,MAAM,GAAG,EAAE,CAAC;QAElB,KAAK,MAAM,MAAM,IAAI,GAAG,CAAC,OAAO,IAAI,EAAE,EAAE,CAAC;YACvC,IAAI,GAAG,CAAC,UAAU,EAAE,CAAC;gBACnB,IAAI,CAAC,KAAK,CAAC,IAAI,CAAC,EAAE,GAAG,SAAS,EAAE,GAAG,GAAG,CAAC,UAAU,EAAE,CAAC,CAAC;YACvD,CAAC;YAED,GAAG,CAAC,SAAS,CAAC,MAAM,EAAE,SAAS,CAAC,CAAC;YAEjC,IAAI,MAAM,CAAC,OAAO,EAAE,CAAC;gBACnB,IAAA,0BAAa,EAAC,IAAI,CAAC,SAAS,EAAE,MAAM,EAAE,MAAM,CAAC,OAAO,CAAC,CAAC;gBACtD,IAAA,6BAAgB,EAAC,IAAI,CAAC,OAAO,EAAE,MAAM,CAAC,CAAC;YACzC,CAAC;YAED,IAAI,MAAM,CAAC,aAAa,EAAE,CAAC;gBACzB,MAAM,KAAK,GAAG,cAAc,CAAC,IAAI,CAAC,EAAE,EAAE,MAAM,CAAC,aAAa,EAAE,MAAM,CAAC,CAAC;gBAEpE,IAAI,KAAK,EAAE,CAAC;oBACV,MAAM,IAAI,CAAC,gBAAgB,CAAC,KAAK,EAAE,GAAG,EAAE,SAAS,EAAE,OAAO,CAAC,CAAC;gBAC9D,CAAC;YACH,CAAC;YAED,IAAI,MAAM,CAAC,YAAY,KAAK,QAAQ,EAAE,CAAC;gBACrC,MAAM,IAAI,KAAK,CAAC,sCAAsC,CAAC,CAAC;YAC1D,CAAC;QACH,CAAC;QAED,OAAO,EAAE,GAAG,MAAM,EAAoB,CAAC;IACzC,CAAC;IAEO,KAAK,CAAC,QAAQ,CACpB,MAAU,EACV,OAA2C,EAC3C,IAAa;QAEb,MAAM,UAAU,GAAG,IAAI,CAAC,OAAO,EAAE,UAAU,IAAI,OAAO,EAAE,UAAU,IAAI,CAAC,CAAC;QACxE,MAAM,QAAQ,GAAG,IAAI,CAAC,OAAO,EAAE,QAAQ,IAAI,OAAO,EAAE,QAAQ,IAAI,EAAE,CAAC;QACnE,MAAM,GAAG,GAAG,IAAI,CAAC,OAAO,EAAE,GAAG,IAAI,OAAO,EAAE,GAAG,IAAI,IAAI,mBAAQ,EAAE,CAAC;QAChE,MAAM,SAAS,GAAG,IAAI,CAAC,EAAE,CAAC,WAAW,EAAE,CAAC,SAAS,CAAC;QAElD,IAAI,GAAmD,CAAC;QAExD,IAAI,IAAI,CAAC,OAAO,KAAK,IAAI,CAAC,SAAS,CAAC,IAAI,EAAE,EAAE,CAAC;YAC3C,MAAM,cAAc,GAAG,IAAI,CAAC,OAAO,EAAE,cAAc,IAAI,4BAAgB,CAAC;YACxE,IAAI,CAAC,EAAE,GAAG,IAAI,cAAc,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC;QAC/C,CAAC;QAED,MAAM,MAAM,GAAG,IAAI,CAAC,EAAE,CAAC,MAAM,CAAK,MAAM,EAAE;YACxC,QAAQ,EAAE,IAAI,CAAC,QAAQ;YACvB,KAAK,EAAE,IAAI,CAAC,KAAK;SAClB,CAAC,CAAC;QAEH,MAAM,OAAO,GAAG,EAAE,IAAI,EAAE,MAAe,EAAE,OAAO,EAAE,MAAM,EAAE,CAAC;QAC3D,GAAG,CAAC,GAAG,CAAC,OAAO,EAAE,OAAO,EAAE,SAAS,CAAC,CAAC;QAErC,aAAa,EAAE,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,QAAQ,EAAE,CAAC,EAAE,EAAE,CAAC;YACjD,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,UAAU,EAAE,CAAC,EAAE,EAAE,CAAC;gBACpC,IAAI,CAAC;oBACH,MAAM,EACJ,SAAS,EACT,OAAO,EACP,WAAW,EACX,MAAM,EAAE,QAAQ,EACjB,GAAG,OAAO,IAAI,EAAE,CAAC;oBAClB,MAAM,MAAM,GAAG,SAAS,IAAI,QAAQ,CAAC;oBAErC,MAAM,MAAM,GAAG,MAAM,IAAI,CAAC,WAAW,CAAC;wBACpC,EAAE,EAAE,OAAO,EAAE,EAAE,IAAI,IAAI,CAAC,EAAE;wBAC1B,GAAG;wBACH,SAAS;wBACT,OAAO;wBACP,WAAW;wBACX,MAAM;wBACN,QAAQ,EAAE,OAAO,EAAE,QAAQ;wBAC3B,WAAW,EAAE,OAAO,EAAE,WAAW;qBAClC,CAAC,CAAC;oBAEH,MAAM,WAAW,GAAG,GAAG,CAAC,OAAO,CAAC,SAAS,CAAC,CAAC;oBAE3C,IAAI,WAAW,EAAE,IAAI,KAAK,UAAU,EAAE,CAAC;wBACrC,SAAS,aAAa,CAAC;oBACzB,CAAC;oBAED,IAAA,iCAAoB,EAAC,IAAI,CAAC,SAAS,EAAE,MAAM,CAAC,CAAC;oBAC7C,IAAI,CAAC,KAAK,GAAG,EAAE,GAAG,MAAM,EAAE,CAAC;oBAC3B,OAAO,MAAM,CAAC;gBAChB,CAAC;gBAAC,OAAO,CAAC,EAAE,CAAC;oBACX,IAAI,WAAW,CAAC;oBAChB,IAAI,EAAE,qBAAqB,CAAC,CAAU,CAAC,CAAC;oBAExC,IAAI,CAAC,YAAY,4BAAe,EAAE,CAAC;wBACjC,WAAW,GAAG,CAAC,CAAC,qBAAqB,EAAE,CAAC;wBACxC,GAAG,GAAG,CAAC,CAAC;oBACV,CAAC;yBAAM,IAAI,CAAC,YAAY,6BAAgB,EAAE,CAAC;wBACzC,MAAM,EAAE,GAAG,CAAqB,CAAC;wBACjC,WAAW,GAAG,EAAE,CAAC,qBAAqB,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC;wBACvD,GAAG,GAAG,CAAC,CAAC;oBACV,CAAC;yBAAM,CAAC;wBACN,MAAM,CAAC,CAAC;oBACV,CAAC;oBAED,IAAI,WAAW,EAAE,CAAC;wBAChB,MAAM,OAAO,GAAG,IAAI,CAAC,EAAE,CAAC,iBAAiB,CAAC,WAAW,CAAC,CAAC;wBACvD,GAAG,CAAC,GAAG,CAAC,EAAE,IAAI,EAAE,MAAe,EAAE,OAAO,EAAE,EAAE,OAAO,EAAE,SAAS,CAAC,CAAC;wBAEhE,IAAI,OAAO,EAAE,KAAK,EAAE,CAAC;4BACnB,OAAO,CAAC,GAAG,CAAC,mBAAmB,EAAE,OAAO,CAAC,CAAC;wBAC5C,CAAC;oBACH,CAAC;gBACH,CAAC;YACH,CAAC;YACD,IAAI,GAAG,YAAY,6BAAgB,IAAI,GAAG,CAAC,WAAW,EAAE,EAAE,CAAC;gBACzD,OAAO,GAAG,CAAC,QAAQ,EAAS,CAAC;YAC/B,CAAC;YAED,MAAM,IAAI,KAAK,CAAC,mCAAmC,GAAG,EAAE,OAAO,EAAE,CAAC,CAAC;QACrE,CAAC;QAED,MAAM,IAAI,KAAK,CAAC,sDAAsD,CAAC,CAAC;IAC1E,CAAC;IAEe,KAAK,CAAC,OAAO,CAC3B,MAAU,EACV,OAA2C;QAE3C,MAAM,MAAM,GAAG,IAAI,CAAC,OAAO,EAAE,MAAM,IAAI,OAAO,EAAE,MAAM,CAAC;QAEvD,IAAI,CAAC,MAAM,EAAE,CAAC;YACZ,OAAO,MAAM,IAAI,CAAC,QAAQ,CAAC,MAAM,EAAE,OAAO,CAAC,CAAC;QAC9C,CAAC;QAED,MAAM,UAAU,GAAG;YACjB,CAAC,oBAAoB,CAAC,EAAE,IAAI,CAAC,SAAS,CAAC,QAAQ,EAAE;YACjD,CAAC,oBAAoB,CAAC,EAAE,IAAI,CAAC,YAAY,IAAI,MAAM;SACpD,CAAC;QAEF,OAAO,MAAM,MAAM,CAAC,eAAe,CACjC,UAAU,EACV;YACE,IAAI,EAAE,qBAAU,CAAC,MAAM;YACvB,UAAU;SACX,EACD,KAAK,EAAE,IAAI,EAAE,EAAE;YACb,MAAM,GAAG,GAAG,IAAI,CAAC,QAAQ,CAAC,MAAM,EAAE,OAAO,EAAE,IAAI,CAAC,CAAC;YACjD,IAAI,CAAC,GAAG,EAAE,CAAC;YACX,OAAO,GAAG,CAAC;QACb,CAAC,CACF,CAAC;IACJ,CAAC;IAEM,gBAAgB,GAAG,KAAK,EAC7B,aAAoD,EACpD,GAAuB,EACvB,SAAkB,EAClB,OAAgB,EAChB,EAAE;QACF,mFAAmF;QACnF,MAAM,QAAQ,GAAG,aAAa,CAAC,GAAG,CAAC,CAAC,IAAI,EAAE,EAAE,CAC1C,IAAI,CAAC,QAAQ,EAAE,OAAO,CAAC,IAAI,EAAE,EAAE,SAAS,EAAE,OAAO,EAAE,CAAC,CAAC,IAAI,CAAC,CAAC,IAAI,EAAE,EAAE;YACjE,IAAI,IAAI,EAAE,EAAE,EAAE,CAAC;gBACb,OAAO;oBACL,IAAI,EAAE,UAAmB;oBACzB,MAAM,EAAE,IAAI,CAAC,MAAM,IAAI,EAAE;oBACzB,UAAU,EAAE,IAAI,CAAC,EAAE;iBACpB,CAAC;YACJ,CAAC;YACD,OAAO,IAAI,CAAC,CAAC,kEAAkE;QACjF,CAAC,CAAC,CACH,CAAC;QAEF,mCAAmC;QACnC,MAAM,OAAO,GAAG,MAAM,OAAO,CAAC,GAAG,CAAC,QAAQ,CAAC,CAAC;QAE5C,OAAO,CAAC,OAAO,CAAC,CAAC,MAAM,EAAE,EAAE;YACzB,IAAI,MAAM,EAAE,CAAC;gBACX,GAAG,CAAC,GAAG,CAAC,MAAM,EAAE,SAAS,CAAC,CAAC;YAC7B,CAAC;QACH,CAAC,CAAC,CAAC;IACL,CAAC,CAAC;CACH;AAjaD,gCAiaC;AAED,SAAS,cAAc,CACrB,EAAyB,EACzB,aAA8D,EAC9D,MAA+B;IAE/B,IAAI,CAAC,aAAa,IAAI,aAAa,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;QACjD,OAAO;IACT,CAAC;IACD,IAAI,EAAE,CAAC,WAAW,EAAE,CAAC,SAAS,EAAE,CAAC;QAC/B,MAAM,KAAK,GAAiC,aAAa,CAAC,GAAG,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,CAAC;YACpE,EAAE,EAAE,CAAC,CAAC,EAAE;YACR,IAAI,EAAE,CAAC,CAAC,QAAQ,CAAC,IAAI;YACrB,IAAI,EAAE,CAAC,CAAC,QAAQ,CAAC,MAAgB;SAClC,CAAC,CAAC,CAAC;QAEJ,0CAA0C;QAC1C,yCAAyC;QACzC,sCAAsC;QACtC,oEAAoE;QACpE,IAAI;QACJ,OAAO,KAAK,CAAC;IACf,CAAC;SAAM,IAAI,MAAM,CAAC,YAAY,EAAE,CAAC;QAC/B,MAAM,EAAE,YAAY,EAAE,iBAAiB,EAAE,GAAG,MAI3C,CAAC;QACF,OAAO,MAAM,CAAC,YAAY,CAAC;QAC3B,OAAO,MAAM,CAAC,iBAAiB,CAAC;QAEhC,OAAO;YACL;gBACE,IAAI,EAAE,YAAY;gBAClB,IAAI,EAAE,iBAAiB;aACxB;SACF,CAAC;IACJ,CAAC;AACH,CAAC"}

package/cjs/dsp/index.d.ts

@@ -0,0 +1,15 @@
+export * from './generate.js';
+export * from './sig.js';
+export * from './optimize.js';
+export * from './evaluate.js';
+export * from './program.js';
+export * from './eval.js';
+export * from './loader.js';
+export * from './strutil.js';
+export * from './router.js';
+export * from './loader.js';
+export type { AxAssertion, AxStreamingAssertion } from './asserts.js';
+export type { AxPromptTemplate, AxFieldTemplateFn } from './prompt.js';
+export type { AxInstanceRegistry } from './registry.js';
+export type { AxResponseHandlerArgs } from './generate.js';
+export type { ParsedField, ParsedIdentifier } from './parser.js';

package/cjs/dsp/index.js

@@ -0,0 +1,27 @@
+"use strict";
+var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    var desc = Object.getOwnPropertyDescriptor(m, k);
+    if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
+      desc = { enumerable: true, get: function() { return m[k]; } };
+    }
+    Object.defineProperty(o, k2, desc);
+}) : (function(o, m, k, k2) {
+    if (k2 === undefined) k2 = k;
+    o[k2] = m[k];
+}));
+var __exportStar = (this && this.__exportStar) || function(m, exports) {
+    for (var p in m) if (p !== "default" && !Object.prototype.hasOwnProperty.call(exports, p)) __createBinding(exports, m, p);
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+__exportStar(require("./generate.js"), exports);
+__exportStar(require("./sig.js"), exports);
+__exportStar(require("./optimize.js"), exports);
+__exportStar(require("./evaluate.js"), exports);
+__exportStar(require("./program.js"), exports);
+__exportStar(require("./eval.js"), exports);
+__exportStar(require("./loader.js"), exports);
+__exportStar(require("./strutil.js"), exports);
+__exportStar(require("./router.js"), exports);
+__exportStar(require("./loader.js"), exports);
+//# sourceMappingURL=index.js.map

package/cjs/dsp/index.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"index.js","sourceRoot":"","sources":["../../../dsp/index.ts"],"names":[],"mappings":";;;;;;;;;;;;;;;;AAAA,gDAA8B;AAC9B,2CAAyB;AACzB,gDAA8B;AAC9B,gDAA8B;AAC9B,+CAA6B;AAC7B,4CAA0B;AAC1B,8CAA4B;AAC5B,+CAA6B;AAC7B,8CAA4B;AAC5B,8CAA4B"}

package/cjs/dsp/loader.d.ts

@@ -0,0 +1,30 @@
+import type { AxFieldValue } from './program.js';
+export type AxDataRow = {
+    row: Record<string, AxFieldValue>;
+};
+export declare class AxHFDataLoader {
+    private rows;
+    private baseUrl;
+    private dataset;
+    private split;
+    private config;
+    private options?;
+    constructor({ dataset, split, config, options }: Readonly<{
+        dataset: string;
+        split: string;
+        config: string;
+        options?: Readonly<{
+            offset?: number;
+            length?: number;
+        }>;
+    }>);
+    private fetchDataFromAPI;
+    loadData(): Promise<AxDataRow[]>;
+    setData(rows: AxDataRow[]): void;
+    getData(): AxDataRow[];
+    getRows<T>({ count, fields, renameMap }: Readonly<{
+        count: number;
+        fields: readonly string[];
+        renameMap?: Record<string, string>;
+    }>): Promise<T[]>;
+}

package/cjs/dsp/loader.js

@@ -0,0 +1,85 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.AxHFDataLoader = void 0;
+class AxHFDataLoader {
+    rows = [];
+    baseUrl;
+    dataset;
+    split;
+    config;
+    options;
+    constructor({ dataset, split, config, options }) {
+        this.baseUrl = 'https://datasets-server.huggingface.co/rows';
+        this.dataset = dataset;
+        this.split = split;
+        this.config = config;
+        this.options = options;
+    }
+    async fetchDataFromAPI(url) {
+        try {
+            const response = await fetch(url);
+            if (!response.ok) {
+                throw new Error(`Error fetching data: ${response.statusText}`);
+            }
+            const data = (await response.json());
+            if (!data?.rows) {
+                throw new Error('Invalid data format');
+            }
+            return data.rows;
+        }
+        catch (error) {
+            console.error('Error fetching data from API:', error);
+            throw error;
+        }
+    }
+    // https://datasets-server.huggingface.co/rows?dataset=hotpot_qa&config=distractor&split=train&offset=0&length=100
+    async loadData() {
+        const offset = this.options?.offset ?? 0;
+        const length = this.options?.length ?? 100;
+        const ds = encodeURIComponent(this.dataset);
+        const url = `${this.baseUrl}?dataset=${ds}&config=${this.config}&split=${this.split}&offset=${offset}&length=${length}`;
+        console.log('Downloading data from API.');
+        this.rows = (await this.fetchDataFromAPI(url));
+        return this.rows;
+    }
+    // eslint-disable-next-line functional/prefer-immutable-types
+    setData(rows) {
+        this.rows = rows;
+    }
+    getData() {
+        return this.rows;
+    }
+    async getRows({ count, fields, renameMap }) {
+        if (this.rows.length === 0) {
+            throw new Error('No data loaded, call loadData or setData first.');
+        }
+        const dataRows = this.rows.slice(0, count);
+        return dataRows
+            .map((item) => {
+            const result = {};
+            fields.forEach((field) => {
+                const keys = field.split('.');
+                // Initial value should match the type of the rows, and be indexable by string
+                let value = item.row;
+                for (const key of keys) {
+                    // Use type assertion to tell TypeScript that value will always be an object that can be indexed with string keys
+                    if (Object.prototype.hasOwnProperty.call(value, key)) {
+                        value = value[key];
+                    }
+                }
+                if (!value) {
+                    return;
+                }
+                const resultFieldName = renameMap && field in renameMap ? renameMap[field] : field;
+                if (!resultFieldName) {
+                    throw new Error(`Invalid field name: ${field}`);
+                }
+                result[resultFieldName] = value;
+            });
+            return result;
+        })
+            .filter((v) => Object.keys(v).length !== 0);
+    }
+}
+exports.AxHFDataLoader = AxHFDataLoader;
+//# sourceMappingURL=loader.js.map

package/cjs/dsp/loader.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"loader.js","sourceRoot":"","sources":["../../../dsp/loader.ts"],"names":[],"mappings":";;;AAIA,MAAa,cAAc;IACjB,IAAI,GAAgB,EAAE,CAAC;IACvB,OAAO,CAAS;IAEhB,OAAO,CAAS;IAChB,KAAK,CAAS;IACd,MAAM,CAAS;IACf,OAAO,CAAkD;IAEjE,YAAY,EACV,OAAO,EACP,KAAK,EACL,MAAM,EACN,OAAO,EAMP;QACA,IAAI,CAAC,OAAO,GAAG,6CAA6C,CAAC;QAC7D,IAAI,CAAC,OAAO,GAAG,OAAO,CAAC;QACvB,IAAI,CAAC,KAAK,GAAG,KAAK,CAAC;QACnB,IAAI,CAAC,MAAM,GAAG,MAAM,CAAC;QACrB,IAAI,CAAC,OAAO,GAAG,OAAO,CAAC;IACzB,CAAC;IAEO,KAAK,CAAC,gBAAgB,CAAC,GAAW;QACxC,IAAI,CAAC;YACH,MAAM,QAAQ,GAAG,MAAM,KAAK,CAAC,GAAG,CAAC,CAAC;YAClC,IAAI,CAAC,QAAQ,CAAC,EAAE,EAAE,CAAC;gBACjB,MAAM,IAAI,KAAK,CAAC,wBAAwB,QAAQ,CAAC,UAAU,EAAE,CAAC,CAAC;YACjE,CAAC;YACD,MAAM,IAAI,GAAG,CAAC,MAAM,QAAQ,CAAC,IAAI,EAAE,CAA0B,CAAC;YAC9D,IAAI,CAAC,IAAI,EAAE,IAAI,EAAE,CAAC;gBAChB,MAAM,IAAI,KAAK,CAAC,qBAAqB,CAAC,CAAC;YACzC,CAAC;YACD,OAAO,IAAI,CAAC,IAAI,CAAC;QACnB,CAAC;QAAC,OAAO,KAAK,EAAE,CAAC;YACf,OAAO,CAAC,KAAK,CAAC,+BAA+B,EAAE,KAAK,CAAC,CAAC;YACtD,MAAM,KAAK,CAAC;QACd,CAAC;IACH,CAAC;IAED,kHAAkH;IAE3G,KAAK,CAAC,QAAQ;QACnB,MAAM,MAAM,GAAG,IAAI,CAAC,OAAO,EAAE,MAAM,IAAI,CAAC,CAAC;QACzC,MAAM,MAAM,GAAG,IAAI,CAAC,OAAO,EAAE,MAAM,IAAI,GAAG,CAAC;QAC3C,MAAM,EAAE,GAAG,kBAAkB,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC;QAE5C,MAAM,GAAG,GAAG,GAAG,IAAI,CAAC,OAAO,YAAY,EAAE,WAAW,IAAI,CAAC,MAAM,UAAU,IAAI,CAAC,KAAK,WAAW,MAAM,WAAW,MAAM,EAAE,CAAC;QAExH,OAAO,CAAC,GAAG,CAAC,4BAA4B,CAAC,CAAC;QAC1C,IAAI,CAAC,IAAI,GAAG,CAAC,MAAM,IAAI,CAAC,gBAAgB,CAAC,GAAG,CAAC,CAAgB,CAAC;QAC9D,OAAO,IAAI,CAAC,IAAI,CAAC;IACnB,CAAC;IAED,6DAA6D;IACtD,OAAO,CAAC,IAAiB;QAC9B,IAAI,CAAC,IAAI,GAAG,IAAI,CAAC;IACnB,CAAC;IAEM,OAAO;QACZ,OAAO,IAAI,CAAC,IAAI,CAAC;IACnB,CAAC;IAEM,KAAK,CAAC,OAAO,CAAI,EACtB,KAAK,EACL,MAAM,EACN,SAAS,EAKT;QACA,IAAI,IAAI,CAAC,IAAI,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;YAC3B,MAAM,IAAI,KAAK,CAAC,iDAAiD,CAAC,CAAC;QACrE,CAAC;QACD,MAAM,QAAQ,GAAG,IAAI,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC,EAAE,KAAK,CAAC,CAAC;QAE3C,OAAO,QAAQ;aACZ,GAAG,CAAC,CAAC,IAAI,EAAE,EAAE;YACZ,MAAM,MAAM,GAAiC,EAAE,CAAC;YAEhD,MAAM,CAAC,OAAO,CAAC,CAAC,KAAK,EAAE,EAAE;gBACvB,MAAM,IAAI,GAAG,KAAK,CAAC,KAAK,CAAC,GAAG,CAAC,CAAC;gBAC9B,8EAA8E;gBAC9E,IAAI,KAAK,GAA2B,IAAI,CAAC,GAAG,CAAC;gBAC7C,KAAK,MAAM,GAAG,IAAI,IAAI,EAAE,CAAC;oBACvB,iHAAiH;oBACjH,IACE,MAAM,CAAC,SAAS,CAAC,cAAc,CAAC,IAAI,CAClC,KAAgC,EAChC,GAAG,CACJ,EACD,CAAC;wBACD,KAAK,GAAI,KAAiC,CAAC,GAAG,CAAC,CAAC;oBAClD,CAAC;gBACH,CAAC;gBACD,IAAI,CAAC,KAAK,EAAE,CAAC;oBACX,OAAO;gBACT,CAAC;gBACD,MAAM,eAAe,GACnB,SAAS,IAAI,KAAK,IAAI,SAAS,CAAC,CAAC,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,KAAK,CAAC;gBAC7D,IAAI,CAAC,eAAe,EAAE,CAAC;oBACrB,MAAM,IAAI,KAAK,CAAC,uBAAuB,KAAK,EAAE,CAAC,CAAC;gBAClD,CAAC;gBACD,MAAM,CAAC,eAAe,CAAC,GAAG,KAAqB,CAAC;YAClD,CAAC,CAAC,CAAC;YAEH,OAAO,MAAM,CAAC;QAChB,CAAC,CAAC;aACD,MAAM,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,MAAM,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC,MAAM,KAAK,CAAC,CAAQ,CAAC;IACvD,CAAC;CACF;AAnHD,wCAmHC"}

package/cjs/dsp/optimize.d.ts

@@ -0,0 +1,27 @@
+import type { AxFieldValue, AxGenIn, AxGenOut, AxProgram, AxProgramDemos } from './program.js';
+export type AxExample = Record<string, AxFieldValue>;
+export type AxMetricFn = <T extends AxGenOut = AxGenOut>(arg0: Readonly<{
+    prediction: T;
+    example: AxExample;
+}>) => boolean;
+export type AxMetricFnArgs = Parameters<AxMetricFn>[0];
+export type AxOptimizerArgs<IN extends AxGenIn, OUT extends AxGenOut> = {
+    program: Readonly<AxProgram<IN, OUT>>;
+    examples: Readonly<AxExample[]>;
+    options?: {
+        maxRounds?: number;
+        maxExamples?: number;
+        maxDemos?: number;
+    };
+};
+export declare class AxBootstrapFewShot<IN extends AxGenIn = AxGenIn, OUT extends AxGenOut = AxGenOut> {
+    private program;
+    private examples;
+    private maxRounds;
+    private maxDemos;
+    private maxExamples;
+    private traces;
+    constructor({ program, examples, options }: Readonly<AxOptimizerArgs<IN, OUT>>);
+    private compileRound;
+    compile(metricFn: AxMetricFn, options?: Readonly<AxOptimizerArgs<IN, OUT>['options']>): Promise<AxProgramDemos[]>;
+}

package/cjs/dsp/optimize.js

@@ -0,0 +1,99 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.AxBootstrapFewShot = void 0;
+const util_js_1 = require("./util.js");
+class AxBootstrapFewShot {
+    program;
+    examples;
+    maxRounds;
+    maxDemos;
+    maxExamples;
+    traces = [];
+    constructor({ program, examples = [], options }) {
+        if (examples.length == 0) {
+            throw new Error('No examples found');
+        }
+        this.maxRounds = options?.maxRounds ?? 3;
+        this.maxDemos = options?.maxDemos ?? 4;
+        this.maxExamples = options?.maxExamples ?? 16;
+        this.program = program;
+        this.examples = examples;
+    }
+    async compileRound(roundIndex, metricFn, options) {
+        const st = new Date().getTime();
+        const maxDemos = options?.maxDemos ?? this.maxDemos;
+        const aiOpt = { modelConfig: { temperature: 0.7 } };
+        const examples = randomSample(this.examples, this.maxExamples);
+        for (let i = 0; i < examples.length; i++) {
+            if (i > 0) {
+                aiOpt.modelConfig.temperature = 0.7 + 0.001 * i;
+            }
+            const ex = examples[i];
+            if (!ex) {
+                throw new Error('Invalid example');
+            }
+            const exList = [...examples.slice(0, i), ...examples.slice(i + 1)];
+            this.program.setExamples(exList);
+            const res = await this.program.forward(ex, aiOpt);
+            const success = metricFn({ prediction: res, example: ex });
+            if (success) {
+                this.traces = [...this.traces, ...this.program.getTraces()];
+            }
+            const current = i + examples.length * roundIndex;
+            const total = examples.length * this.maxRounds;
+            const et = new Date().getTime() - st;
+            (0, util_js_1.updateProgressBar)(current, total, this.traces.length, et, 30, 'Tuning Prompt');
+            if (this.traces.length > maxDemos) {
+                return;
+            }
+        }
+    }
+    async compile(metricFn, options) {
+        const maxRounds = options?.maxRounds ?? this.maxRounds;
+        this.traces = [];
+        for (let i = 0; i < maxRounds; i++) {
+            await this.compileRound(i, metricFn, options);
+        }
+        if (this.traces.length === 0) {
+            throw new Error('No demonstrations found. Either provider more examples or improve the existing ones.');
+        }
+        const demos = groupTracesByKeys(this.traces);
+        return demos;
+    }
+}
+exports.AxBootstrapFewShot = AxBootstrapFewShot;
+function groupTracesByKeys(programTraces) {
+    const groupedTraces = new Map();
+    // Group all traces by their keys
+    for (const programTrace of programTraces) {
+        if (groupedTraces.has(programTrace.programId)) {
+            groupedTraces.get(programTrace.programId).push(programTrace.trace);
+        }
+        else {
+            groupedTraces.set(programTrace.programId, [programTrace.trace]);
+        }
+    }
+    // Convert the Map into an array of ProgramDemos
+    const programDemosArray = [];
+    groupedTraces.forEach((traces, programId) => {
+        programDemosArray.push({ traces, programId });
+    });
+    return programDemosArray;
+}
+const randomSample = (array, n) => {
+    // Clone the array to avoid modifying the original array
+    const clonedArray = [...array];
+    // Shuffle the cloned array
+    for (let i = clonedArray.length - 1; i > 0; i--) {
+        const j = Math.floor(Math.random() * (i + 1));
+        const caI = clonedArray[i];
+        const caJ = clonedArray[j];
+        if (!caI || !caJ) {
+            throw new Error('Invalid array elements');
+        }
+        [clonedArray[i], clonedArray[j]] = [caJ, caI];
+    }
+    // Return the first `n` items of the shuffled array
+    return clonedArray.slice(0, n);
+};
+//# sourceMappingURL=optimize.js.map

package/cjs/dsp/optimize.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"optimize.js","sourceRoot":"","sources":["../../../dsp/optimize.ts"],"names":[],"mappings":";;;AAQA,uCAA8C;AAgB9C,MAAa,kBAAkB;IAIrB,OAAO,CAA+B;IACtC,QAAQ,CAAwB;IAChC,SAAS,CAAS;IAClB,QAAQ,CAAS;IACjB,WAAW,CAAS;IACpB,MAAM,GAAqB,EAAE,CAAC;IAEtC,YAAY,EACV,OAAO,EACP,QAAQ,GAAG,EAAE,EACb,OAAO,EAC4B;QACnC,IAAI,QAAQ,CAAC,MAAM,IAAI,CAAC,EAAE,CAAC;YACzB,MAAM,IAAI,KAAK,CAAC,mBAAmB,CAAC,CAAC;QACvC,CAAC;QACD,IAAI,CAAC,SAAS,GAAG,OAAO,EAAE,SAAS,IAAI,CAAC,CAAC;QACzC,IAAI,CAAC,QAAQ,GAAG,OAAO,EAAE,QAAQ,IAAI,CAAC,CAAC;QACvC,IAAI,CAAC,WAAW,GAAG,OAAO,EAAE,WAAW,IAAI,EAAE,CAAC;QAE9C,IAAI,CAAC,OAAO,GAAG,OAAO,CAAC;QACvB,IAAI,CAAC,QAAQ,GAAG,QAAQ,CAAC;IAC3B,CAAC;IAEO,KAAK,CAAC,YAAY,CACxB,UAAkB,EAClB,QAAoB,EACpB,OAAuD;QAEvD,MAAM,EAAE,GAAG,IAAI,IAAI,EAAE,CAAC,OAAO,EAAE,CAAC;QAChC,MAAM,QAAQ,GAAG,OAAO,EAAE,QAAQ,IAAI,IAAI,CAAC,QAAQ,CAAC;QACpD,MAAM,KAAK,GAAG,EAAE,WAAW,EAAE,EAAE,WAAW,EAAE,GAAG,EAAE,EAAE,CAAC;QACpD,MAAM,QAAQ,GAAG,YAAY,CAAC,IAAI,CAAC,QAAQ,EAAE,IAAI,CAAC,WAAW,CAAC,CAAC;QAE/D,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,QAAQ,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;YACzC,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC;gBACV,KAAK,CAAC,WAAW,CAAC,WAAW,GAAG,GAAG,GAAG,KAAK,GAAG,CAAC,CAAC;YAClD,CAAC;YAED,MAAM,EAAE,GAAG,QAAQ,CAAC,CAAC,CAAC,CAAC;YACvB,IAAI,CAAC,EAAE,EAAE,CAAC;gBACR,MAAM,IAAI,KAAK,CAAC,iBAAiB,CAAC,CAAC;YACrC,CAAC;YACD,MAAM,MAAM,GAAG,CAAC,GAAG,QAAQ,CAAC,KAAK,CAAC,CAAC,EAAE,CAAC,CAAC,EAAE,GAAG,QAAQ,CAAC,KAAK,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC;YACnE,IAAI,CAAC,OAAO,CAAC,WAAW,CAAC,MAAM,CAAC,CAAC;YAEjC,MAAM,GAAG,GAAG,MAAM,IAAI,CAAC,OAAO,CAAC,OAAO,CAAC,EAAQ,EAAE,KAAK,CAAC,CAAC;YACxD,MAAM,OAAO,GAAG,QAAQ,CAAC,EAAE,UAAU,EAAE,GAAG,EAAE,OAAO,EAAE,EAAE,EAAE,CAAC,CAAC;YAC3D,IAAI,OAAO,EAAE,CAAC;gBACZ,IAAI,CAAC,MAAM,GAAG,CAAC,GAAG,IAAI,CAAC,MAAM,EAAE,GAAG,IAAI,CAAC,OAAO,CAAC,SAAS,EAAE,CAAC,CAAC;YAC9D,CAAC;YAED,MAAM,OAAO,GAAG,CAAC,GAAG,QAAQ,CAAC,MAAM,GAAG,UAAU,CAAC;YACjD,MAAM,KAAK,GAAG,QAAQ,CAAC,MAAM,GAAG,IAAI,CAAC,SAAS,CAAC;YAC/C,MAAM,EAAE,GAAG,IAAI,IAAI,EAAE,CAAC,OAAO,EAAE,GAAG,EAAE,CAAC;YAErC,IAAA,2BAAiB,EACf,OAAO,EACP,KAAK,EACL,IAAI,CAAC,MAAM,CAAC,MAAM,EAClB,EAAE,EACF,EAAE,EACF,eAAe,CAChB,CAAC;YAEF,IAAI,IAAI,CAAC,MAAM,CAAC,MAAM,GAAG,QAAQ,EAAE,CAAC;gBAClC,OAAO;YACT,CAAC;QACH,CAAC;IACH,CAAC;IAEM,KAAK,CAAC,OAAO,CAClB,QAAoB,EACpB,OAAuD;QAEvD,MAAM,SAAS,GAAG,OAAO,EAAE,SAAS,IAAI,IAAI,CAAC,SAAS,CAAC;QACvD,IAAI,CAAC,MAAM,GAAG,EAAE,CAAC;QAEjB,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,SAAS,EAAE,CAAC,EAAE,EAAE,CAAC;YACnC,MAAM,IAAI,CAAC,YAAY,CAAC,CAAC,EAAE,QAAQ,EAAE,OAAO,CAAC,CAAC;QAChD,CAAC;QAED,IAAI,IAAI,CAAC,MAAM,CAAC,MAAM,KAAK,CAAC,EAAE,CAAC;YAC7B,MAAM,IAAI,KAAK,CACb,sFAAsF,CACvF,CAAC;QACJ,CAAC;QAED,MAAM,KAAK,GAAqB,iBAAiB,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;QAC/D,OAAO,KAAK,CAAC;IACf,CAAC;CACF;AA9FD,gDA8FC;AAED,SAAS,iBAAiB,CACxB,aAAwC;IAExC,MAAM,aAAa,GAAG,IAAI,GAAG,EAA0C,CAAC;IAExE,iCAAiC;IACjC,KAAK,MAAM,YAAY,IAAI,aAAa,EAAE,CAAC;QACzC,IAAI,aAAa,CAAC,GAAG,CAAC,YAAY,CAAC,SAAS,CAAC,EAAE,CAAC;YAC9C,aAAa,CAAC,GAAG,CAAC,YAAY,CAAC,SAAS,CAAE,CAAC,IAAI,CAAC,YAAY,CAAC,KAAK,CAAC,CAAC;QACtE,CAAC;aAAM,CAAC;YACN,aAAa,CAAC,GAAG,CAAC,YAAY,CAAC,SAAS,EAAE,CAAC,YAAY,CAAC,KAAK,CAAC,CAAC,CAAC;QAClE,CAAC;IACH,CAAC;IAED,gDAAgD;IAChD,MAAM,iBAAiB,GAAqB,EAAE,CAAC;IAC/C,aAAa,CAAC,OAAO,CAAC,CAAC,MAAM,EAAE,SAAS,EAAE,EAAE;QAC1C,iBAAiB,CAAC,IAAI,CAAC,EAAE,MAAM,EAAE,SAAS,EAAE,CAAC,CAAC;IAChD,CAAC,CAAC,CAAC;IAEH,OAAO,iBAAiB,CAAC;AAC3B,CAAC;AAED,MAAM,YAAY,GAAG,CAAI,KAAmB,EAAE,CAAS,EAAO,EAAE;IAC9D,wDAAwD;IACxD,MAAM,WAAW,GAAG,CAAC,GAAG,KAAK,CAAC,CAAC;IAC/B,2BAA2B;IAC3B,KAAK,IAAI,CAAC,GAAG,WAAW,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC;QAChD,MAAM,CAAC,GAAG,IAAI,CAAC,KAAK,CAAC,IAAI,CAAC,MAAM,EAAE,GAAG,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC;QAC9C,MAAM,GAAG,GAAG,WAAW,CAAC,CAAC,CAAC,CAAC;QAC3B,MAAM,GAAG,GAAG,WAAW,CAAC,CAAC,CAAC,CAAC;QAE3B,IAAI,CAAC,GAAG,IAAI,CAAC,GAAG,EAAE,CAAC;YACjB,MAAM,IAAI,KAAK,CAAC,wBAAwB,CAAC,CAAC;QAC5C,CAAC;QAED,CAAC,WAAW,CAAC,CAAC,CAAC,EAAE,WAAW,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,GAAG,EAAE,GAAG,CAAC,CAAC;IAChD,CAAC;IACD,mDAAmD;IACnD,OAAO,WAAW,CAAC,KAAK,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC;AACjC,CAAC,CAAC"}