@amitdeshmukh/ax-crew 7.0.0 → 8.0.0

package/dist/agents/ace.js

@@ -0,0 +1,477 @@
+/**
+ * ACE (Agentic Context Engineering) integration for AxCrew
+ *
+ * This module provides helpers to build and manage AxACE optimizers for agents,
+ * enabling offline compilation and online learning from feedback.
+ *
+ * Reference: https://axllm.dev/ace/
+ */
+import { AxACE, ai as buildAI, AxSignature, AxGen } from "@ax-llm/ax";
+/**
+ * Create an empty playbook structure
+ */
+export const createEmptyPlaybook = () => {
+    const now = new Date().toISOString();
+    return {
+        version: 1,
+        sections: {},
+        stats: {
+            bulletCount: 0,
+            helpfulCount: 0,
+            harmfulCount: 0,
+            tokenEstimate: 0,
+        },
+        updatedAt: now,
+    };
+};
+/**
+ * Render a playbook into markdown instruction block for injection into prompts.
+ * Mirrors the AxACE renderPlaybook function.
+ */
+export const renderPlaybook = (playbook) => {
+    if (!playbook)
+        return '';
+    const sectionsObj = playbook.sections || {};
+    const header = playbook.description
+        ? `## Context Playbook\n${playbook.description.trim()}\n`
+        : '## Context Playbook\n';
+    const sectionEntries = Object.entries(sectionsObj);
+    if (sectionEntries.length === 0)
+        return '';
+    const sections = sectionEntries
+        .map(([sectionName, bullets]) => {
+        const body = bullets
+            .map((bullet) => `- [${bullet.id}] ${bullet.content}`)
+            .join('\n');
+        return body
+            ? `### ${sectionName}\n${body}`
+            : `### ${sectionName}\n_(empty)_`;
+    })
+        .join('\n\n');
+    return `${header}\n${sections}`.trim();
+};
+/**
+ * Check if running in Node.js environment (for file operations)
+ */
+const isNodeLike = () => {
+    try {
+        return typeof process !== "undefined" && !!process.versions?.node;
+    }
+    catch {
+        return false;
+    }
+};
+/**
+ * Read JSON file (Node.js only)
+ */
+const readFileJSON = async (path) => {
+    if (!isNodeLike())
+        return undefined;
+    try {
+        const { readFile } = await import("fs/promises");
+        const buf = await readFile(path, "utf-8");
+        return JSON.parse(buf);
+    }
+    catch {
+        return undefined;
+    }
+};
+/**
+ * Write JSON file (Node.js only)
+ */
+const writeFileJSON = async (path, data) => {
+    if (!isNodeLike())
+        return;
+    try {
+        const { mkdir, writeFile } = await import("fs/promises");
+        const { dirname } = await import("path");
+        await mkdir(dirname(path), { recursive: true });
+        await writeFile(path, JSON.stringify(data ?? {}, null, 2), "utf-8");
+    }
+    catch {
+        // Swallow persistence errors by default
+    }
+};
+/**
+ * Resolve environment variable
+ */
+const resolveEnv = (name) => {
+    try {
+        if (typeof process !== "undefined" && process?.env) {
+            return process.env[name];
+        }
+        return globalThis?.[name];
+    }
+    catch {
+        return undefined;
+    }
+};
+/**
+ * Build teacher AI instance from config, falling back to student AI
+ */
+const buildTeacherAI = (teacherCfg, fallback) => {
+    if (!teacherCfg)
+        return fallback;
+    const { provider, providerKeyName, apiURL, ai: aiConfig, providerArgs } = teacherCfg;
+    if (!provider || !providerKeyName || !aiConfig)
+        return fallback;
+    const apiKey = resolveEnv(providerKeyName) || "";
+    if (!apiKey)
+        return fallback;
+    const args = {
+        name: provider,
+        apiKey,
+        config: aiConfig,
+        options: {}
+    };
+    if (apiURL)
+        args.apiURL = apiURL;
+    if (providerArgs && typeof providerArgs === "object") {
+        Object.assign(args, providerArgs);
+    }
+    try {
+        return buildAI(args);
+    }
+    catch {
+        return fallback;
+    }
+};
+/**
+ * Build an AxACE optimizer for an agent
+ *
+ * @param studentAI - The agent's AI instance (used as student)
+ * @param cfg - ACE configuration
+ * @returns Configured AxACE optimizer
+ */
+export const buildACEOptimizer = (studentAI, cfg) => {
+    const teacherAI = buildTeacherAI(cfg.teacher, studentAI);
+    // Build optimizer options, only include initialPlaybook if it has the right structure
+    const optimizerOptions = {
+        maxEpochs: cfg.options?.maxEpochs,
+        allowDynamicSections: cfg.options?.allowDynamicSections,
+    };
+    // Only pass initialPlaybook if it looks like a valid playbook structure
+    if (cfg.persistence?.initialPlaybook &&
+        typeof cfg.persistence.initialPlaybook === 'object' &&
+        'sections' in cfg.persistence.initialPlaybook) {
+        optimizerOptions.initialPlaybook = cfg.persistence.initialPlaybook;
+    }
+    return new AxACE({
+        studentAI,
+        teacherAI,
+        verbose: !!cfg.options?.maxEpochs
+    }, optimizerOptions);
+};
+/**
+ * Load initial playbook from file, callback, or inline config
+ *
+ * @param cfg - Persistence configuration
+ * @returns Loaded playbook or undefined
+ */
+export const loadInitialPlaybook = async (cfg) => {
+    if (!cfg)
+        return undefined;
+    // Try callback first
+    if (typeof cfg.onLoad === "function") {
+        try {
+            return await cfg.onLoad();
+        }
+        catch {
+            // Fall through to other methods
+        }
+    }
+    // Try inline playbook
+    if (cfg.initialPlaybook) {
+        return cfg.initialPlaybook;
+    }
+    // Try file path
+    if (cfg.playbookPath) {
+        return await readFileJSON(cfg.playbookPath);
+    }
+    return undefined;
+};
+/**
+ * Persist playbook to file or via callback
+ *
+ * @param pb - Playbook to persist
+ * @param cfg - Persistence configuration
+ */
+export const persistPlaybook = async (pb, cfg) => {
+    if (!cfg || !pb)
+        return;
+    // Call persist callback if provided
+    if (typeof cfg.onPersist === "function") {
+        try {
+            await cfg.onPersist(pb);
+        }
+        catch {
+            // Ignore callback errors
+        }
+    }
+    // Write to file if auto-persist enabled
+    if (cfg.autoPersist && cfg.playbookPath) {
+        await writeFileJSON(cfg.playbookPath, pb);
+    }
+};
+/**
+ * Resolve metric function from registry or create equality-based metric
+ *
+ * @param cfg - Metric configuration
+ * @param registry - Function registry to search
+ * @returns Metric function or undefined
+ */
+export const resolveMetric = (cfg, registry) => {
+    if (!cfg)
+        return undefined;
+    const { metricFnName, primaryOutputField } = cfg;
+    // Try to find a function by name in the registry
+    if (metricFnName) {
+        const candidate = registry[metricFnName];
+        if (typeof candidate === "function") {
+            return candidate;
+        }
+    }
+    // Create simple equality-based metric if primary output field specified
+    if (primaryOutputField) {
+        const field = primaryOutputField;
+        return ({ prediction, example }) => {
+            try {
+                return prediction?.[field] === example?.[field] ? 1 : 0;
+            }
+            catch {
+                return 0;
+            }
+        };
+    }
+    return undefined;
+};
+/**
+ * Run offline ACE compilation
+ *
+ * @param args - Compilation arguments
+ * @returns Compilation result with optimized program
+ */
+export const runOfflineCompile = async (args) => {
+    const { program, optimizer, metric, examples = [], persistence } = args;
+    if (!optimizer || !metric || examples.length === 0) {
+        return null;
+    }
+    try {
+        // Run compilation
+        const result = await optimizer.compile(program, examples, metric);
+        // Extract and persist playbook
+        const playbook = result?.artifact?.playbook;
+        if (playbook && persistence) {
+            await persistPlaybook(playbook, persistence);
+        }
+        return result;
+    }
+    catch (error) {
+        console.warn("ACE offline compile failed:", error);
+        return null;
+    }
+};
+/**
+ * Apply online update with feedback
+ *
+ * @param args - Update arguments
+ * @returns Curator delta (operations applied)
+ */
+export const runOnlineUpdate = async (args) => {
+    const { optimizer, example, prediction, feedback, persistence, debug } = args;
+    if (!optimizer)
+        return null;
+    try {
+        // Apply online update (per ACE API: example, prediction, feedback)
+        const curatorDelta = await optimizer.applyOnlineUpdate({
+            example,
+            prediction,
+            feedback
+        });
+        // Access the optimizer's private playbook property
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        const playbook = optimizer.playbook;
+        // Persist updated playbook if we have one and persistence is configured
+        if (playbook && persistence?.autoPersist) {
+            await persistPlaybook(playbook, persistence);
+        }
+        return curatorDelta;
+    }
+    catch (error) {
+        // AxACE's reflector sometimes returns bulletTags in non-array format, causing iteration errors.
+        // This is a known issue - we fall back to direct playbook updates via addFeedbackToPlaybook.
+        if (debug) {
+            console.warn("[ACE Debug] AxACE applyOnlineUpdate failed (falling back to direct update):", error);
+        }
+        return null;
+    }
+};
+/**
+ * Generate a unique bullet ID (mirrors AxACE's generateBulletId)
+ */
+const generateBulletId = (section) => {
+    const normalized = section
+        .toLowerCase()
+        .replace(/[^a-z0-9]+/g, '-')
+        .replace(/^-+|-+$/g, '')
+        .slice(0, 6);
+    const randomHex = Math.random().toString(16).slice(2, 10);
+    return `${normalized || 'ctx'}-${randomHex}`;
+};
+/**
+ * Recompute playbook stats after modifications
+ */
+const recomputePlaybookStats = (playbook) => {
+    let bulletCount = 0;
+    let helpfulCount = 0;
+    let harmfulCount = 0;
+    let tokenEstimate = 0;
+    const sections = playbook.sections || {};
+    for (const bullets of Object.values(sections)) {
+        for (const bullet of bullets) {
+            bulletCount += 1;
+            helpfulCount += bullet.helpfulCount;
+            harmfulCount += bullet.harmfulCount;
+            tokenEstimate += Math.ceil(bullet.content.length / 4);
+        }
+    }
+    playbook.stats = { bulletCount, helpfulCount, harmfulCount, tokenEstimate };
+    playbook.updatedAt = new Date().toISOString();
+};
+/**
+ * Apply curator operations to playbook (mirrors AxACE's applyCuratorOperations)
+ */
+const applyCuratorOperations = (playbook, operations) => {
+    // Ensure playbook has sections initialized
+    if (!playbook.sections) {
+        playbook.sections = {};
+    }
+    const now = new Date().toISOString();
+    for (const op of operations) {
+        if (!op.section)
+            continue;
+        // Initialize section if needed
+        if (!playbook.sections[op.section]) {
+            playbook.sections[op.section] = [];
+        }
+        const section = playbook.sections[op.section];
+        switch (op.type) {
+            case 'ADD': {
+                if (!op.content?.trim())
+                    continue;
+                // Check for duplicates
+                const isDuplicate = section.some(b => b.content.toLowerCase() === op.content.toLowerCase());
+                if (isDuplicate)
+                    continue;
+                const bullet = {
+                    id: op.bulletId || generateBulletId(op.section),
+                    section: op.section,
+                    content: op.content.trim(),
+                    helpfulCount: 1,
+                    harmfulCount: 0,
+                    createdAt: now,
+                    updatedAt: now,
+                };
+                section.push(bullet);
+                break;
+            }
+            case 'UPDATE': {
+                if (!op.bulletId)
+                    continue;
+                const bullet = section.find(b => b.id === op.bulletId);
+                if (bullet && op.content) {
+                    bullet.content = op.content.trim();
+                    bullet.updatedAt = now;
+                }
+                break;
+            }
+            case 'REMOVE': {
+                if (!op.bulletId)
+                    continue;
+                const idx = section.findIndex(b => b.id === op.bulletId);
+                if (idx >= 0)
+                    section.splice(idx, 1);
+                break;
+            }
+        }
+    }
+    recomputePlaybookStats(playbook);
+};
+// Cached feedback analyzer program (created lazily)
+let feedbackAnalyzerProgram = null;
+/**
+ * Get or create the feedback analyzer program.
+ * Uses AxGen with a proper signature, just like AxACE's reflector/curator.
+ *
+ * Uses `class` type for section to get type-safe enums and better token efficiency.
+ * See: https://axllm.dev/signatures/
+ */
+const getOrCreateFeedbackAnalyzer = () => {
+    if (!feedbackAnalyzerProgram) {
+        const signature = new AxSignature(`feedback:string "User feedback to analyze"
+       -> 
+       section:class "Guidelines, Response Strategies, Common Pitfalls, Root Cause Notes" "Playbook section category",
+       content:string "The specific instruction to add to the playbook - keep all concrete details"`);
+        signature.setDescription(`Convert user feedback into a playbook instruction. Keep ALL specific details from the feedback (times, names, numbers, constraints).`);
+        feedbackAnalyzerProgram = new AxGen(signature);
+    }
+    return feedbackAnalyzerProgram;
+};
+/**
+ * Use LLM to analyze feedback and generate playbook operations.
+ *
+ * This leverages AxGen with a proper signature (like AxACE's reflector/curator)
+ * to properly categorize feedback and extract actionable insights.
+ *
+ * IMPORTANT: The prompt explicitly tells the LLM to preserve specificity.
+ *
+ * @param ai - The AI instance to use for analysis
+ * @param feedback - User feedback string
+ * @param debug - Whether to log debug info
+ * @returns Promise of curator operations
+ */
+export const analyzeAndCategorizeFeedback = async (ai, feedback, debug = false) => {
+    if (!feedback?.trim())
+        return [];
+    try {
+        const analyzer = getOrCreateFeedbackAnalyzer();
+        const result = await analyzer.forward(ai, {
+            feedback: feedback.trim(),
+        });
+        if (debug) {
+            console.log('[ACE Debug] Feedback analysis result:', result);
+        }
+        // Section is guaranteed to be valid by the class type constraint
+        const section = result.section || 'Guidelines';
+        // Use the LLM's content, but fall back to raw feedback if empty
+        const content = result.content?.trim() || feedback.trim();
+        return [{ type: 'ADD', section, content }];
+    }
+    catch (error) {
+        if (debug) {
+            console.warn('[ACE Debug] Feedback analysis failed, using raw feedback:', error);
+        }
+        // Fallback: use the raw feedback as-is
+        return [{ type: 'ADD', section: 'Guidelines', content: feedback.trim() }];
+    }
+};
+/**
+ * Add feedback to playbook using LLM analysis.
+ *
+ * Uses the AI to properly understand and categorize the feedback,
+ * then applies it as a curator operation.
+ *
+ * @param playbook - The playbook to update (mutated in place)
+ * @param feedback - User feedback string to add
+ * @param ai - AI instance for smart categorization
+ * @param debug - Whether to log debug info
+ */
+export const addFeedbackToPlaybook = async (playbook, feedback, ai, debug = false) => {
+    if (!playbook || !feedback?.trim())
+        return;
+    // Use LLM to categorize feedback while preserving specificity
+    const operations = await analyzeAndCategorizeFeedback(ai, feedback, debug);
+    if (operations.length > 0) {
+        applyCuratorOperations(playbook, operations);
+    }
+};

package/dist/agents/agentConfig.d.ts

@@ -35,5 +35,6 @@ declare const parseAgentConfig: (agentName: string, crewConfig: AxCrewConfig, fu
     subAgentNames: string[];
     examples: Record<string, any>[];
     tracker: AxDefaultCostTracker;
+    debug: any;
 }>;
 export { parseAgentConfig, parseCrewConfig };

package/dist/agents/agentConfig.js

@@ -179,6 +179,7 @@ const parseAgentConfig = async (agentName, crewConfig, functions, state, options
             subAgentNames: agentConfigData.agents || [],
             examples: agentConfigData.examples || [],
             tracker: costTracker,
+            debug: agentConfigData.options?.debug ?? agentConfigData.debug ?? false,
         };
     }
     catch (error) {

package/dist/agents/index.d.ts

@@ -1,12 +1,17 @@
 import { AxAgent, AxAI } from "@ax-llm/ax";
 import type { AxSignature, AxAgentic, AxFunction, AxProgramForwardOptions, AxProgramStreamingForwardOptions, AxGenStreamingOut } from "@ax-llm/ax";
-import type { StateInstance, FunctionRegistryType, UsageCost, AxCrewConfig, AxCrewOptions, MCPTransportConfig } from "../types.js";
+import type { StateInstance, FunctionRegistryType, UsageCost, AxCrewConfig, AxCrewOptions, MCPTransportConfig, ACEConfig } from "../types.js";
 declare class StatefulAxAgent extends AxAgent<any, any> {
     state: StateInstance;
     axai: any;
     private agentName;
     private costTracker?;
     private lastRecordedCostUSD;
+    private debugEnabled;
+    private aceConfig?;
+    private aceOptimizer?;
+    private acePlaybook?;
+    private aceBaseInstruction?;
     private isAxAIService;
     private isAxAIInstance;
     constructor(ai: AxAI, options: Readonly<{
@@ -18,6 +23,7 @@ declare class StatefulAxAgent extends AxAgent<any, any> {
         functions?: (AxFunction | (() => AxFunction))[] | undefined;
         examples?: Array<Record<string, any>> | undefined;
         mcpServers?: Record<string, MCPTransportConfig> | undefined;
+        debug?: boolean;
     }>, state: StateInstance);
     forward(values: Record<string, any>, options?: Readonly<AxProgramForwardOptions<any>>): Promise<Record<string, any>>;
     forward(ai: AxAI, values: Record<string, any>, options?: Readonly<AxProgramForwardOptions<any>>): Promise<Record<string, any>>;
@@ -37,6 +43,51 @@ declare class StatefulAxAgent extends AxAgent<any, any> {
      * Call this to start fresh measurement windows for the agent.
      */
     resetMetrics(): void;
+    /**
+     * Initialize ACE (Agentic Context Engineering) for this agent.
+     * Builds the optimizer and loads any initial playbook from persistence.
+     * Sets up the optimizer for online-only mode if compileOnStart is false.
+     */
+    initACE(ace?: ACEConfig): Promise<void>;
+    /**
+     * Run offline ACE compilation with examples and metric.
+     * Compiles the playbook based on training examples.
+     */
+    optimizeOffline(params?: {
+        metric?: any;
+        examples?: any[];
+    }): Promise<void>;
+    /**
+     * Apply online ACE update based on user feedback.
+     *
+     * For preference-based feedback (e.g., "only show flights between 9am-12pm"),
+     * we use our own feedback analyzer that preserves specificity.
+     *
+     * Note: AxACE's built-in curator is designed for error correction (severity mismatches)
+     * and tends to over-abstract preference feedback into generic guidelines.
+     * We bypass it and directly use our feedback analyzer for better results.
+     */
+    applyOnlineUpdate(params: {
+        example: any;
+        prediction: any;
+        feedback?: string;
+    }): Promise<void>;
+    /**
+     * Get the current ACE playbook for this agent.
+     */
+    getPlaybook(): any | undefined;
+    /**
+     * Apply an ACE playbook to this agent.
+     * Stores the playbook for use in next forward() call.
+     * Note: Playbook is composed into instruction BEFORE each forward(), mirroring AxACE.compile behavior.
+     */
+    applyPlaybook(pb: any): void;
+    /**
+     * Compose instruction with current playbook and set on agent.
+     * This mirrors what AxACE does internally before each forward() during compile().
+     * Should be called BEFORE forward() to ensure playbook is in the prompt.
+     */
+    private composeInstructionWithPlaybook;
 }
 /**
  * AxCrew orchestrates a set of Ax agents that share state,
@@ -64,6 +115,7 @@ declare class AxCrew {
     crewId: string;
     agents: Map<string, StatefulAxAgent> | null;
     state: StateInstance;
+    private executionHistory;
     /**
      * Creates an instance of AxCrew.
      * @param {AxCrewConfig} crewConfig - JSON object with crew configuration.
@@ -93,6 +145,36 @@ declare class AxCrew {
      */
     addAgentsToCrew(agentNames: string[]): Promise<Map<string, StatefulAxAgent> | null>;
     addAllAgents(): Promise<Map<string, StatefulAxAgent> | null>;
+    /**
+     * Track agent execution for ACE feedback routing
+     */
+    trackAgentExecution(taskId: string, agentName: string, input: any): void;
+    /**
+     * Record agent result for ACE feedback routing
+     */
+    recordAgentResult(taskId: string, agentName: string, result: any): void;
+    /**
+     * Get agent involvement for a task (used for ACE feedback routing)
+     */
+    getTaskAgentInvolvement(taskId: string): {
+        rootAgent: string;
+        involvedAgents: string[];
+        taskInput: any;
+        agentResults: Map<string, any>;
+        duration?: number;
+    } | null;
+    /**
+     * Apply feedback to agents involved in a task for ACE online learning
+     */
+    applyTaskFeedback(params: {
+        taskId: string;
+        feedback: string;
+        strategy?: 'all' | 'primary' | 'weighted';
+    }): Promise<void>;
+    /**
+     * Clean up old execution history (call periodically to prevent memory leaks)
+     */
+    cleanupOldExecutions(maxAgeMs?: number): void;
     /**
      * Cleans up the crew by dereferencing agents and resetting the state.
      */