@amitdeshmukh/ax-crew 7.0.0 → 8.0.0

package/dist/agents/index.js

@@ -10,6 +10,13 @@ class StatefulAxAgent extends AxAgent {
     agentName;
     costTracker;
     lastRecordedCostUSD = 0;
+    debugEnabled = false;
+    // ACE-related optional state
+    aceConfig;
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    aceOptimizer;
+    acePlaybook;
+    aceBaseInstruction; // Original description before playbook injection
     isAxAIService(obj) {
         return !!obj && typeof obj.getName === 'function' && typeof obj.chat === 'function';
     }
@@ -17,7 +24,7 @@ class StatefulAxAgent extends AxAgent {
         return !!obj && typeof obj === 'object' && ('defaults' in obj || 'modelInfo' in obj);
     }
     constructor(ai, options, state) {
-        const { examples, ...restOptions } = options;
+        const { examples, debug, ...restOptions } = options;
         const formattedOptions = {
             ...restOptions,
             functions: restOptions.functions?.map((fn) => typeof fn === "function" ? fn() : fn),
@@ -26,6 +33,7 @@ class StatefulAxAgent extends AxAgent {
         this.state = state;
         this.axai = ai;
         this.agentName = options.name;
+        this.debugEnabled = debug ?? false;
         // Set examples if provided
         if (examples && examples.length > 0) {
             super.setExamples(examples);
@@ -37,8 +45,34 @@ class StatefulAxAgent extends AxAgent {
         const start = performance.now();
         const crewId = this.state?.crewId || (this.state.get?.('crewId')) || 'default';
         const labels = { crewId, agent: this.agentName };
-        // Track costs regardless of whether it's a direct or sub-agent call
-        // This ensures we capture multiple legitimate calls to the same agent
+        const input = this.isAxAIService(first) ? second : first;
+        const taskId = `task_${crewId}_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;
+        // Track execution context in crew for ACE feedback routing
+        const crewInstance = this.state?.crew;
+        if (crewInstance) {
+            if (this.isAxAIService(first)) {
+                // For sub-agent calls, track under parent task ID
+                const parentTaskId = this.state?.currentTaskId;
+                if (parentTaskId) {
+                    crewInstance.trackAgentExecution(parentTaskId, this.agentName, input);
+                }
+            }
+            else {
+                // Root-level call - start new execution tracking
+                crewInstance.trackAgentExecution(taskId, this.agentName, input);
+                this.state.currentTaskId = taskId;
+            }
+        }
+        // Before forward: compose instruction with current playbook (mirrors AxACE.compile behavior)
+        // This ensures the agent uses the latest playbook context for this call
+        if (this.debugEnabled) {
+            console.log(`[ACE Debug] forward() called, aceConfig=${!!this.aceConfig}`);
+        }
+        if (this.aceConfig) {
+            await this.composeInstructionWithPlaybook();
+        }
+        // Execute the forward call
+        // Note: OpenTelemetry spans are automatically created by AxAI (configured via AxCrewOptions.telemetry)
         if (this.isAxAIService(first)) {
             // Sub-agent case (called with AI service)
             result = await super.forward(this.axai, second, third);
@@ -83,6 +117,24 @@ class StatefulAxAgent extends AxAgent {
             }
             catch { }
         }
+        // Record result in crew execution history for ACE feedback routing
+        if (crewInstance) {
+            if (this.isAxAIService(first)) {
+                // For sub-agent calls, record under parent task ID
+                const parentTaskId = this.state?.currentTaskId;
+                if (parentTaskId) {
+                    crewInstance.recordAgentResult(parentTaskId, this.agentName, result);
+                }
+            }
+            else {
+                // Root-level result - include taskId for feedback routing
+                crewInstance.recordAgentResult(taskId, this.agentName, result);
+                // Clean up current task ID
+                delete this.state.currentTaskId;
+                // Attach taskId to result for feedback routing convenience
+                result._taskId = taskId;
+            }
+        }
         return result;
     }
     // Implementation
@@ -174,6 +226,196 @@ class StatefulAxAgent extends AxAgent {
         const crewId = this.state?.crewId || (this.state.get?.('crewId')) || 'default';
         MetricsRegistry.reset({ crewId, agent: this.agentName });
     }
+    // =============
+    // ACE API - Agentic Context Engineering for online learning
+    // Reference: https://axllm.dev/ace/
+    // =============
+    /**
+     * 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.
+     */
+    async initACE(ace) {
+        this.aceConfig = ace;
+        if (!ace)
+            return;
+        try {
+            // Capture base instruction BEFORE any playbook injection (mirrors AxACE.extractProgramInstruction)
+            this.aceBaseInstruction = this.getSignature().getDescription() || '';
+            const { buildACEOptimizer, loadInitialPlaybook, createEmptyPlaybook } = await import('./ace.js');
+            // Build optimizer with agent's AI as student
+            this.aceOptimizer = buildACEOptimizer(this.axai, ace);
+            // For online-only mode (no offline compile), we need to set the program
+            // reference so applyOnlineUpdate can work. AxACE requires compile() to
+            // set the program, but we can set it directly for online-only use.
+            if (!ace.compileOnStart) {
+                // eslint-disable-next-line @typescript-eslint/no-explicit-any
+                this.aceOptimizer.program = this;
+            }
+            // Load initial playbook or create empty one
+            const initial = await loadInitialPlaybook(ace.persistence);
+            this.applyPlaybook(initial ?? createEmptyPlaybook());
+            if (this.debugEnabled) {
+                console.log(`[ACE Debug] Initialized for ${this.agentName}, base instruction: ${this.aceBaseInstruction?.slice(0, 50)}...`);
+            }
+        }
+        catch (error) {
+            console.warn(`Failed to initialize ACE for agent ${this.agentName}:`, error);
+        }
+    }
+    /**
+     * Run offline ACE compilation with examples and metric.
+     * Compiles the playbook based on training examples.
+     */
+    async optimizeOffline(params) {
+        if (!this.aceConfig || !this.aceOptimizer)
+            return;
+        try {
+            const { runOfflineCompile, resolveMetric } = await import('./ace.js');
+            const registry = this.__functionsRegistry;
+            const metric = params?.metric || resolveMetric(this.aceConfig.metric, registry || {});
+            const examples = params?.examples || [];
+            if (!metric || examples.length === 0) {
+                console.warn(`ACE offline compile skipped for ${this.agentName}: missing metric or examples`);
+                return;
+            }
+            const result = await runOfflineCompile({
+                program: this,
+                optimizer: this.aceOptimizer,
+                metric,
+                examples,
+                persistence: this.aceConfig.persistence
+            });
+            // Apply optimized playbook if compilation succeeded
+            if (result?.artifact?.playbook) {
+                await this.applyPlaybook(result.artifact.playbook);
+            }
+        }
+        catch (error) {
+            console.warn(`ACE offline compile failed for ${this.agentName}:`, error);
+        }
+    }
+    /**
+     * 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.
+     */
+    async applyOnlineUpdate(params) {
+        if (!this.aceConfig)
+            return;
+        if (!params.feedback?.trim())
+            return; // Nothing to do without feedback
+        try {
+            const { persistPlaybook, addFeedbackToPlaybook, createEmptyPlaybook } = await import('./ace.js');
+            // Get or create playbook
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            let playbook = this.acePlaybook ?? this.aceOptimizer?.playbook;
+            if (!playbook) {
+                playbook = createEmptyPlaybook();
+            }
+            if (this.debugEnabled) {
+                console.log(`[ACE Debug] Adding feedback to playbook: "${params.feedback}"`);
+            }
+            // Use teacher AI (or student AI as fallback) for smart categorization
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            const teacherAI = this.aceOptimizer?.teacherAI;
+            const aiForAnalysis = teacherAI ?? this.axai;
+            // Directly add feedback to playbook using our analyzer (preserves specificity)
+            await addFeedbackToPlaybook(playbook, params.feedback, aiForAnalysis, this.debugEnabled);
+            // Store updated playbook (injection happens in next forward() call)
+            this.applyPlaybook(playbook);
+            // Sync with optimizer if available
+            if (this.aceOptimizer) {
+                this.aceOptimizer.playbook = playbook;
+            }
+            // Persist if auto-persist enabled
+            if (this.aceConfig.persistence?.autoPersist) {
+                await persistPlaybook(playbook, this.aceConfig.persistence);
+            }
+            if (this.debugEnabled) {
+                console.log(`[ACE Debug] Playbook updated, sections: ${Object.keys(playbook.sections || {}).join(', ')}`);
+            }
+        }
+        catch (error) {
+            console.warn(`ACE online update failed for ${this.agentName}:`, error);
+        }
+    }
+    /**
+     * Get the current ACE playbook for this agent.
+     */
+    getPlaybook() {
+        return this.acePlaybook;
+    }
+    /**
+     * 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) {
+        this.acePlaybook = pb;
+        // Also update optimizer's internal playbook if possible
+        try {
+            this.aceOptimizer.playbook = pb;
+        }
+        catch {
+            // Ignore - optimizer may not be initialized yet
+        }
+    }
+    /**
+     * 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.
+     */
+    async composeInstructionWithPlaybook() {
+        const playbook = this.acePlaybook ?? this.aceOptimizer?.playbook;
+        if (this.debugEnabled) {
+            console.log(`[ACE Debug] composeInstructionWithPlaybook called`);
+            console.log(`[ACE Debug] playbook exists: ${!!playbook}, sections: ${playbook ? Object.keys(playbook.sections || {}).length : 0}`);
+            console.log(`[ACE Debug] baseInstruction: "${this.aceBaseInstruction?.slice(0, 50)}..."`);
+        }
+        if (!playbook)
+            return;
+        try {
+            const { renderPlaybook } = await import('./ace.js');
+            const rendered = renderPlaybook(playbook);
+            if (this.debugEnabled) {
+                console.log(`[ACE Debug] rendered playbook (${rendered.length} chars): ${rendered.slice(0, 100)}...`);
+            }
+            if (!rendered)
+                return;
+            // Compose: base instruction + playbook (just like AxACE.composeInstruction)
+            const baseInstruction = this.aceBaseInstruction || '';
+            const combinedInstruction = [baseInstruction.trim(), '', rendered]
+                .filter((part) => part.trim().length > 0)
+                .join('\n\n');
+            if (this.debugEnabled) {
+                console.log(`[ACE Debug] combinedInstruction (${combinedInstruction.length} chars)`);
+            }
+            if (combinedInstruction.length >= 20) {
+                // Call setDescription on the internal program (like AxACE does)
+                // AxAgent.setDescription() only updates the signature, but we need
+                // to update the program's description which is used for the system prompt
+                const program = this.program;
+                if (program?.setDescription) {
+                    program.setDescription(combinedInstruction);
+                }
+                // Also update via AxAgent's setDescription for consistency
+                this.setDescription(combinedInstruction);
+                if (this.debugEnabled) {
+                    console.log(`[ACE Debug] setDescription called successfully`);
+                    console.log(`[ACE Debug] Verifying - signature desc length: ${this.getSignature().getDescription()?.length}`);
+                }
+            }
+        }
+        catch (error) {
+            console.warn('[ACE Debug] Failed to compose instruction:', error);
+        }
+    }
 }
 /**
  * AxCrew orchestrates a set of Ax agents that share state,
@@ -201,6 +443,8 @@ class AxCrew {
     crewId;
     agents;
     state;
+    // Execution history for ACE feedback routing
+    executionHistory = new Map();
     /**
      * Creates an instance of AxCrew.
      * @param {AxCrewConfig} crewConfig - JSON object with crew configuration.
@@ -270,6 +514,8 @@ class AxCrew {
                 }
             }
             // Create an instance of StatefulAxAgent
+            // Set crew reference in state for execution tracking (ACE feedback routing)
+            const agentState = { ...this.state, crew: this };
             const agent = new StatefulAxAgent(ai, {
                 name,
                 description,
@@ -278,8 +524,24 @@ class AxCrew {
                 functions: uniqueFunctions,
                 agents: uniqueSubAgents,
                 examples,
-            }, this.state);
+                debug: agentConfig.debug,
+            }, agentState);
             agent.costTracker = tracker;
+            agent.__functionsRegistry = this.functionsRegistry;
+            // Initialize ACE if configured
+            try {
+                const crewAgent = parseCrewConfig(this.crewConfig).crew.find(a => a.name === name);
+                const ace = crewAgent?.ace;
+                if (ace) {
+                    await agent.initACE?.(ace);
+                    if (ace.compileOnStart) {
+                        const { resolveMetric } = await import('./ace.js');
+                        const metric = resolveMetric(ace.metric, this.functionsRegistry);
+                        await agent.optimizeOffline?.({ metric, examples });
+                    }
+                }
+            }
+            catch { }
             return agent;
         }
         catch (error) {
@@ -397,11 +659,104 @@ class AxCrew {
             throw error;
         }
     }
+    /**
+     * Track agent execution for ACE feedback routing
+     */
+    trackAgentExecution(taskId, agentName, input) {
+        if (!this.executionHistory.has(taskId)) {
+            this.executionHistory.set(taskId, {
+                taskId,
+                rootAgent: agentName,
+                involvedAgents: new Set([agentName]),
+                taskInput: input,
+                agentResults: new Map(),
+                startTime: Date.now()
+            });
+        }
+        else {
+            // Add to involved agents if not already present
+            const context = this.executionHistory.get(taskId);
+            context.involvedAgents.add(agentName);
+        }
+    }
+    /**
+     * Record agent result for ACE feedback routing
+     */
+    recordAgentResult(taskId, agentName, result) {
+        const context = this.executionHistory.get(taskId);
+        if (context) {
+            context.agentResults.set(agentName, result);
+            context.endTime = Date.now();
+        }
+    }
+    /**
+     * Get agent involvement for a task (used for ACE feedback routing)
+     */
+    getTaskAgentInvolvement(taskId) {
+        const context = this.executionHistory.get(taskId);
+        if (!context)
+            return null;
+        return {
+            rootAgent: context.rootAgent,
+            involvedAgents: Array.from(context.involvedAgents),
+            taskInput: context.taskInput,
+            agentResults: context.agentResults,
+            duration: context.endTime ? context.endTime - context.startTime : undefined
+        };
+    }
+    /**
+     * Apply feedback to agents involved in a task for ACE online learning
+     */
+    async applyTaskFeedback(params) {
+        const involvement = this.getTaskAgentInvolvement(params.taskId);
+        if (!involvement) {
+            console.warn(`No execution history found for task ${params.taskId}`);
+            return;
+        }
+        const { involvedAgents, taskInput, agentResults } = involvement;
+        const strategy = params.strategy || 'all';
+        // Determine which agents to update based on strategy
+        let agentsToUpdate = [];
+        if (strategy === 'primary') {
+            agentsToUpdate = [involvement.rootAgent];
+        }
+        else if (strategy === 'all' || strategy === 'weighted') {
+            agentsToUpdate = involvedAgents;
+        }
+        // Apply feedback to each involved agent
+        for (const agentName of agentsToUpdate) {
+            const agent = this.agents?.get(agentName);
+            if (agent && typeof agent.applyOnlineUpdate === 'function') {
+                try {
+                    await agent.applyOnlineUpdate({
+                        example: taskInput,
+                        prediction: agentResults.get(agentName),
+                        feedback: params.feedback
+                    });
+                }
+                catch (error) {
+                    console.warn(`Failed to apply ACE feedback to agent ${agentName}:`, error);
+                }
+            }
+        }
+    }
+    /**
+     * Clean up old execution history (call periodically to prevent memory leaks)
+     */
+    cleanupOldExecutions(maxAgeMs = 3600000) {
+        const cutoffTime = Date.now() - maxAgeMs;
+        for (const [taskId, context] of this.executionHistory) {
+            if (context.startTime < cutoffTime) {
+                this.executionHistory.delete(taskId);
+            }
+        }
+    }
     /**
      * Cleans up the crew by dereferencing agents and resetting the state.
      */
     destroy() {
         this.agents = null;
+        this.executionHistory.clear();
         this.state.reset();
     }
     /**

package/dist/index.d.ts

@@ -1,7 +1,7 @@
 import { AxCrew } from './agents/index.js';
 import { AxCrewFunctions } from './functions/index.js';
-import type { AxCrewConfig, AgentConfig } from './types.js';
-import type { UsageCost, AggregatedMetrics, AggregatedCosts, StateInstance, FunctionRegistryType } from './types.js';
+import type { AxCrewConfig, AxCrewOptions, AgentConfig } from './types.js';
+import type { UsageCost, AggregatedMetrics, AggregatedCosts, StateInstance, FunctionRegistryType, ACEConfig, ACETeacherConfig, ACEPersistenceConfig, ACEOptionsConfig, ACEMetricConfig } from './types.js';
 /**
  * Metrics types and helpers for request counts, token usage, and estimated cost.
  *
@@ -29,4 +29,4 @@ export {
 /** See class JSDoc on the `AxCrew` implementation. */
 _AxCrew as AxCrew, 
 /** Built-in function registry; see file docs in `src/functions/index.ts`. */
-_AxCrewFunctions as AxCrewFunctions, FunctionRegistryType, type AggregatedMetrics, type AggregatedCosts, type AgentConfig, type AxCrewConfig, type StateInstance, type UsageCost, };
+_AxCrewFunctions as AxCrewFunctions, FunctionRegistryType, type AggregatedMetrics, type AggregatedCosts, type AgentConfig, type AxCrewConfig, type AxCrewOptions, type StateInstance, type UsageCost, type ACEConfig, type ACETeacherConfig, type ACEPersistenceConfig, type ACEOptionsConfig, type ACEMetricConfig, };

package/dist/types.d.ts

@@ -127,6 +127,40 @@ interface MCPStreamableHTTPTransportConfig {
  * @property {MCPStdioTransportConfig | MCPHTTPSSETransportConfig | MCPStreambleHTTPTransportConfig} config - The config for the MCP server. Config can be either stdio, http-sse, or streamable http transport.
  */
 type MCPTransportConfig = MCPStdioTransportConfig | MCPHTTPSSETransportConfig | MCPStreamableHTTPTransportConfig;
+interface ACETeacherConfig {
+    provider?: Provider;
+    providerKeyName?: string;
+    apiURL?: string;
+    ai?: AxModelConfig & {
+        model: string;
+    };
+    providerArgs?: Record<string, unknown>;
+}
+interface ACEPersistenceConfig {
+    playbookPath?: string;
+    initialPlaybook?: Record<string, any>;
+    autoPersist?: boolean;
+    onPersist?: (pb: any) => Promise<void> | void;
+    onLoad?: () => Promise<any> | any;
+}
+interface ACEOptionsConfig {
+    maxEpochs?: number;
+    allowDynamicSections?: boolean;
+    tokenBudget?: number;
+    reflectorPrompt?: string;
+    curatorPrompt?: string;
+}
+interface ACEMetricConfig {
+    metricFnName?: string;
+    primaryOutputField?: string;
+}
+interface ACEConfig {
+    teacher?: ACETeacherConfig;
+    persistence?: ACEPersistenceConfig;
+    options?: ACEOptionsConfig;
+    metric?: ACEMetricConfig;
+    compileOnStart?: boolean;
+}
 /**
  * The configuration for an agent.
  *
@@ -175,6 +209,8 @@ interface AgentConfig {
     agents?: string[];
     examples?: Array<Record<string, any>>;
     mcpServers?: Record<string, MCPTransportConfig>;
+    /** Optional AxACE configuration to enable optimization for this agent */
+    ace?: ACEConfig;
 }
 /**
  * The configuration object for an AxCrew instance.
@@ -224,9 +260,11 @@ interface AxCrewConfig {
  * @property {any} [telemetry.meter] - OpenTelemetry Meter instance.
  */
 interface AxCrewOptions {
+    /** Enable debug logging for ACE and other internal operations */
+    debug?: boolean;
     telemetry?: {
         tracer?: any;
         meter?: any;
     };
 }
-export { type AgentConfig, type AxCrewConfig, type AxCrewOptions, type AggregatedMetrics, type StateInstance, type FunctionRegistryType, type MCPStdioTransportConfig, type MCPHTTPSSETransportConfig, type MCPStreamableHTTPTransportConfig, type MCPTransportConfig, type ModelUsage, type ModelInfo, type UsageCost, type AggregatedCosts };
+export { type AgentConfig, type AxCrewConfig, type AxCrewOptions, type AggregatedMetrics, type StateInstance, type FunctionRegistryType, type MCPStdioTransportConfig, type MCPHTTPSSETransportConfig, type MCPStreamableHTTPTransportConfig, type MCPTransportConfig, type ModelUsage, type ModelInfo, type UsageCost, type AggregatedCosts, type ACEConfig, type ACEMetricConfig, type ACEOptionsConfig, type ACEPersistenceConfig, type ACETeacherConfig };

package/examples/README.md

@@ -122,18 +122,56 @@ Key features:
 - Integration with external APIs via MCP servers
 - Cost tracking across agents
 
-### 5. AxCrew Basic Usage
-[`axcrew.ts`](./axcrew.ts)
+### 5. ACE Flight Assistant (Feedback Learning)
+[`ace-feedback-routing.ts`](./ace-feedback-routing.ts)
 
-A basic example showing core AxCrew features:
-- Setting up a crew
-- Adding agents
-- Using shared state
-- Cost tracking
-- Error handling
+Demonstrates ACE (Agentic Context Engineering) learning from user feedback:
+- A flight booking assistant using Google Flights MCP server
+- Real-time playbook updates based on user feedback
+- Persistence of learned preferences across sessions
+- Interactive CLI for flight queries
+
+**Required Setup:**
+```bash
+# Clone the Google Flights MCP server
+git clone https://github.com/opspawn/Google-Flights-MCP-Server.git
+cd Google-Flights-MCP-Server
+python -m venv .venv && source .venv/bin/activate
+pip install -r requirements.txt && playwright install
+
+# Set environment variable
+export GOOGLE_FLIGHTS_SERVER_PATH=/path/to/server.py
+```
+
+Key features:
+- ACE feedback loop integration
+- Playbook persistence to JSON
+- Learning from natural language feedback
+
+### 6. ACE Customer Support (Edge Case Learning)
+[`ace-customer-support.ts`](./ace-customer-support.ts)
+
+Shows how ACE learns to handle edge cases beyond standard policies:
+- A customer support agent with predefined company policies
+- Learning exceptions from supervisor feedback (e.g., "loyal customers get extended returns")
+- Applying learned exceptions to similar future cases
+- Interactive ticket handling demo
 
 **Required Dependencies:** None beyond core package
 
+```typescript
+// Example feedback that teaches the agent:
+// "For loyal customers (5+ years), extend return window to 60 days"
+// "Medical emergencies justify extended deadlines"
+// "Defective products override 'final sale' policy"
+```
+
+Key features:
+- Edge case learning from feedback
+- Policy exception handling
+- Playbook persistence showing learned rules
+- Sample tickets demonstrating various scenarios
+
 ## Running the Examples
 
 1. Clone the repository: