@amitdeshmukh/ax-crew 8.7.3 → 9.0.0

package/dist/agents/deferredTools.js

@@ -0,0 +1,237 @@
+const DEFAULT_THRESHOLD = 20;
+const DEFAULT_MAX_RESULTS = 10;
+/**
+ * Manages deferred tool loading for agents with many tools.
+ *
+ * When tool count exceeds a threshold, only core tools + a `search_tools`
+ * meta-function are visible to the LLM. The LLM calls `search_tools` to
+ * discover and activate deferred tools, which are injected via ax-llm
+ * step hooks for subsequent turns.
+ */
+export class DeferredToolManager {
+    registry;
+    deferredNames;
+    activatedNames;
+    coreNames;
+    maxSearchResults;
+    _isActive;
+    resourceCache;
+    constructor(allFunctions, mcpFunctionNames, config) {
+        const threshold = config?.threshold ?? DEFAULT_THRESHOLD;
+        this.maxSearchResults = config?.maxSearchResults ?? DEFAULT_MAX_RESULTS;
+        this.activatedNames = new Set();
+        this.resourceCache = new Map();
+        // Build full registry
+        this.registry = new Map();
+        for (const fn of allFunctions) {
+            this.registry.set(fn.name, fn);
+        }
+        // Determine if we should activate deferred mode
+        const explicitEnabled = config?.enabled;
+        this._isActive = explicitEnabled !== undefined
+            ? explicitEnabled
+            : allFunctions.length > threshold;
+        // Classify core vs deferred
+        const explicitCore = new Set(config?.coreTools ?? []);
+        this.coreNames = new Set();
+        this.deferredNames = new Set();
+        if (this._isActive) {
+            for (const fn of allFunctions) {
+                const isMcp = mcpFunctionNames.has(fn.name);
+                const isResource = fn.name.startsWith('resource_');
+                const isExplicitCore = explicitCore.has(fn.name);
+                if (isExplicitCore || isResource || !isMcp) {
+                    this.coreNames.add(fn.name);
+                }
+                else {
+                    this.deferredNames.add(fn.name);
+                }
+            }
+        }
+        else {
+            for (const fn of allFunctions) {
+                this.coreNames.add(fn.name);
+            }
+        }
+    }
+    /** Whether deferred mode is active */
+    get isActive() {
+        return this._isActive;
+    }
+    /** Get the initial function set (core tools + search_tools) */
+    getInitialFunctions() {
+        const initial = [];
+        for (const name of this.coreNames) {
+            const fn = this.registry.get(name);
+            if (fn)
+                initial.push(fn.name.startsWith('resource_') ? this.wrapWithCache(fn) : fn);
+        }
+        if (this._isActive) {
+            initial.push(this.createSearchToolFunction());
+        }
+        return initial;
+    }
+    /** Wrap a resource function so repeated calls return cached results */
+    wrapWithCache(fn) {
+        const cache = this.resourceCache;
+        const originalFunc = fn.func;
+        if (!originalFunc)
+            return fn;
+        return {
+            ...fn,
+            func: async (args) => {
+                const cached = cache.get(fn.name);
+                if (cached !== undefined)
+                    return cached;
+                const result = await originalFunc(args);
+                cache.set(fn.name, result);
+                return result;
+            },
+        };
+    }
+    /**
+     * Get step hooks for dynamic tool activation.
+     * - beforeStep: re-injects previously activated tools at the start of each
+     *   forward() call so tools discovered in earlier calls persist.
+     * - afterFunctionExecution: injects newly discovered tools after search_tools
+     *   runs, and auto-activates tools mentioned in function results.
+     */
+    getStepHooks() {
+        const injectedNames = new Set();
+        const injectActivated = (ctx) => {
+            const toInject = [];
+            for (const name of this.activatedNames) {
+                if (!this.coreNames.has(name) && !injectedNames.has(name)) {
+                    const fn = this.registry.get(name);
+                    if (fn) {
+                        toInject.push(fn);
+                        injectedNames.add(name);
+                    }
+                }
+            }
+            if (toInject.length > 0) {
+                ctx.addFunctions(toInject);
+            }
+        };
+        return {
+            beforeStep: async (ctx) => {
+                if (this.activatedNames.size > 0) {
+                    injectActivated(ctx);
+                }
+            },
+            afterFunctionExecution: async (ctx) => {
+                if (ctx.functionsExecuted.has('search_tools')) {
+                    injectActivated(ctx);
+                }
+                // Auto-activate deferred tools mentioned in function results.
+                // Handles cases where a tool error suggests using another tool
+                // (e.g., GraphJin's "recommended_tool": "fix_query_error").
+                this.autoActivateFromResults(ctx.lastFunctionCalls);
+                injectActivated(ctx);
+            },
+        };
+    }
+    /** Scan function results for mentions of deferred tool names and auto-activate them. */
+    autoActivateFromResults(functionCalls) {
+        if (!functionCalls || functionCalls.length === 0 || this.deferredNames.size === 0)
+            return;
+        for (const call of functionCalls) {
+            const resultText = typeof call.result === 'string'
+                ? call.result
+                : JSON.stringify(call.result ?? '');
+            if (!resultText)
+                continue;
+            for (const name of this.deferredNames) {
+                if (!this.activatedNames.has(name) && resultText.includes(name)) {
+                    this.activatedNames.add(name);
+                }
+            }
+        }
+    }
+    /** Search deferred tools by keyword matching on name + description */
+    search(query) {
+        const queryLower = query.toLowerCase();
+        const terms = queryLower.split(/[\s_\-./]+/).filter(t => t.length > 1);
+        const scored = [];
+        for (const name of this.deferredNames) {
+            const fn = this.registry.get(name);
+            if (!fn)
+                continue;
+            const searchText = `${fn.name} ${fn.description ?? ''}`.toLowerCase();
+            let score = 0;
+            for (const term of terms) {
+                if (fn.name.toLowerCase().includes(term))
+                    score += 3;
+                if ((fn.description ?? '').toLowerCase().includes(term))
+                    score += 1;
+            }
+            if (searchText.includes(queryLower))
+                score += 2;
+            if (score > 0) {
+                scored.push({ name, score });
+            }
+        }
+        scored.sort((a, b) => b.score - a.score);
+        const matchedNames = scored.slice(0, this.maxSearchResults).map(s => s.name);
+        if (matchedNames.length === 0) {
+            const names = Array.from(this.deferredNames).slice(0, 30);
+            return `No tools found matching "${query}". Available tools: ${names.join(', ')}${this.deferredNames.size > 30 ? '...' : ''}`;
+        }
+        // Separate newly activated from already-active
+        const newlyActivated = [];
+        const alreadyActive = [];
+        for (const name of matchedNames) {
+            if (this.activatedNames.has(name)) {
+                alreadyActive.push(name);
+            }
+            else {
+                this.activatedNames.add(name);
+                const fn = this.registry.get(name);
+                if (fn)
+                    newlyActivated.push(fn);
+            }
+        }
+        if (newlyActivated.length === 0) {
+            return `All ${alreadyActive.length} matching tools are already active: ${alreadyActive.join(', ')}. Call them directly.`;
+        }
+        const lines = newlyActivated.map((fn) => {
+            const params = fn.parameters?.properties
+                ? Object.keys(fn.parameters.properties).join(', ')
+                : 'none';
+            return `- **${fn.name}**: ${fn.description ?? 'No description'} (params: ${params})`;
+        });
+        const parts = [
+            `Found ${newlyActivated.length} new tool(s) matching "${query}":`,
+            '',
+            ...lines,
+            '',
+            'These tools are now available. Call them directly.',
+        ];
+        if (alreadyActive.length > 0) {
+            parts.push(`Also already active: ${alreadyActive.join(', ')}`);
+        }
+        return parts.join('\n');
+    }
+    /** Create the search_tools meta-function */
+    createSearchToolFunction() {
+        return {
+            name: 'search_tools',
+            description: 'Search for available tools by describing what you need. This agent has additional specialized tools not shown by default. ' +
+                'Describe the task (e.g., "query database tables" or "list available schemas") and matching tools will be activated. ' +
+                'Call this ONCE to discover tools, then use them directly. Do NOT call search_tools again for already discovered tools.',
+            parameters: {
+                type: 'object',
+                properties: {
+                    query: {
+                        type: 'string',
+                        description: 'Describe what you need to do (e.g., "query database", "list tables", "execute graphql")',
+                    },
+                },
+                required: ['query'],
+            },
+            func: async (args) => {
+                return this.search(args.query);
+            },
+        };
+    }
+}

package/dist/agents/index.d.ts

@@ -1,266 +1,4 @@
-import { AxAgent, AxAI, AxSignature as AxSignatureClass } from "@ax-llm/ax";
-import type { AxSignature, AxAgentic, AxFunction, AxProgramForwardOptions, AxProgramStreamingForwardOptions, AxGenStreamingOut } from "@ax-llm/ax";
-import type { StateInstance, FunctionRegistryType, UsageCost, AxCrewConfig, AxCrewOptions, MCPTransportConfig, ACEConfig, AgentExecutionMode, AxCrewAxAgentOptions } from "../types.js";
-declare class StatefulAxAgent extends AxAgent<any, any> {
-    state: StateInstance;
-    axai: any;
-    private agentName;
-    private agentDefinition;
-    private executionMode;
-    private axGenProgram;
-    private costTracker?;
-    private debugEnabled;
-    private static readonly modernAxAgentRuntime;
-    private aceConfig?;
-    private aceOptimizer?;
-    private acePlaybook?;
-    private aceBaseInstruction?;
-    private isAxAIService;
-    constructor(ai: AxAI, options: Readonly<{
-        name: string;
-        description: string;
-        executionMode?: AgentExecutionMode;
-        axAgentOptions?: AxCrewAxAgentOptions;
-        definition?: string;
-        signature: string | AxSignature;
-        agents?: AxAgentic<any, any>[] | undefined;
-        functions?: (AxFunction | (() => AxFunction))[] | undefined;
-        examples?: Array<Record<string, any>> | undefined;
-        mcpServers?: Record<string, MCPTransportConfig> | undefined;
-        debug?: boolean;
-    }>, state: StateInstance);
-    /**
-     * @deprecated Use setExamplesCompat() to avoid Ax runtime version coupling.
-     */
-    setExamples(examples: Readonly<Array<Record<string, any>>>): void;
-    setExamplesCompat(examples: Readonly<Array<Record<string, any>>>): void;
-    /**
-     * @deprecated Use setDescriptionCompat() to avoid Ax runtime version coupling.
-     */
-    setDescription(description: string): void;
-    setDescriptionCompat(description: string): void;
-    getUsage(): (import("@ax-llm/ax").AxModelUsage & {
-        ai: string;
-        model: string;
-    })[];
-    resetUsage(): void;
-    private resolveInvocationArgs;
-    private executeForwardByMode;
-    private recordUsageMetrics;
-    private runForwardInvocation;
-    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>>;
-    private runStreamingInvocation;
-    streamingForward(values: Record<string, any>, options?: Readonly<AxProgramStreamingForwardOptions<any>>): AxGenStreamingOut<any>;
-    streamingForward(ai: AxAI, values: Record<string, any>, options?: Readonly<AxProgramStreamingForwardOptions<any>>): AxGenStreamingOut<any>;
-    getLastUsageCost(): UsageCost | null;
-    getAccumulatedCosts(): UsageCost | null;
-    /**
-     * Get the current metrics snapshot for this agent.
-     * Includes request counts, error rates, token usage, estimated USD cost, and function call stats.
-     *
-     * @returns A metrics snapshot scoped to this agent within its crew.
-     */
-    getMetrics(): import("../metrics/types.js").MetricsSnapshot;
-    /**
-     * Reset all tracked metrics for this agent (does not affect other agents).
-     * 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;
-}
-/**
- * Lightweight proxy that stands in for a real agent in the crew's agent map.
- * It exposes the same `getFunction()` interface (built from the crew config)
- * but defers the expensive `createAgent()` call — and therefore MCP server
- * startup — until the Manager actually delegates to it.
- *
- * Usage: `crew.addLazyAgent("CreateChart")` instead of `crew.addAgent("CreateChart")`
- */
-declare class LazyStatefulAxAgent {
-    private realAgent;
-    private crewRef;
-    private agentName;
-    private description;
-    private signatureStr;
-    private func;
-    private _id;
-    constructor(crewRef: any, agentName: string, crewConfig: AxCrewConfig);
-    private resolve;
-    getFunction(): AxFunction;
-    getSignature(): AxSignatureClass<Record<string, any>, Record<string, any>>;
-    getId(): string;
-    setId(id: string): void;
-    getTraces(): any[];
-    setDemos(): void;
-    getUsage(): any[];
-    resetUsage(): void;
-    forward(...args: any[]): Promise<any>;
-    streamingForward(...args: any[]): any;
-}
-/**
- * AxCrew orchestrates a set of Ax agents that share state,
- * tools (functions), optional MCP servers, streaming, and a built-in metrics
- * registry for tokens, requests, and estimated cost.
- *
- * Typical usage:
- *  const crew = new AxCrew(config, AxCrewFunctions)
- *  await crew.addAllAgents()
- *  const planner = crew.agents?.get("Planner")
- *  const res = await planner?.forward({ task: "Plan something" })
- *
- * Key behaviors:
- * - Validates and instantiates agents from a config-first model
- * - Shares a mutable state object across all agents in the crew
- * - Supports sub-agents and a function registry per agent
- * - Tracks per-agent and crew-level metrics via MetricsRegistry
- * - Provides helpers to add agents (individually, a subset, or all) and
- *   to reset metrics/costs when needed
- */
-declare class AxCrew {
-    private crewConfig;
-    private options?;
-    functionsRegistry: FunctionRegistryType;
-    crewId: string;
-    agents: Map<string, StatefulAxAgent> | null;
-    state: StateInstance;
-    private executionHistory;
-    /**
-     * Creates an instance of AxCrew.
-     * @param {AxCrewConfig} crewConfig - JSON object with crew configuration.
-     * @param {FunctionRegistryType} [functionsRegistry={}] - The registry of functions to use in the crew.
-     * @param {AxCrewOptions} [options] - Optional settings for the crew (e.g., telemetry).
-     * @param {string} [crewId=uuidv4()] - The unique identifier for the crew.
-     */
-    constructor(crewConfig: AxCrewConfig, functionsRegistry?: FunctionRegistryType, options?: AxCrewOptions, crewId?: string);
-    /**
-     * Factory function for creating an agent.
-     * @param {string} agentName - The name of the agent to create.
-     * @returns {StatefulAxAgent} The created StatefulAxAgent instance.
-     * @throws Will throw an error if the agent creation fails.
-     */
-    createAgent: (agentName: string) => Promise<StatefulAxAgent>;
-    /**
-     * Adds an agent to the crew by name.
-     * @param {string} agentName - The name of the agent to add.
-     */
-    addAgent(agentName: string): Promise<void>;
-    /**
-     * Adds a lazy agent to the crew by name.
-     * The agent's tool schema is built immediately from the crew config,
-     * but the expensive initialization (MCP servers, AI client, etc.) is
-     * deferred until the Manager actually delegates to this agent.
-     *
-     * Use this for sub-agents that may not be needed on every request
-     * (e.g., agents with stdio MCP servers that spawn a process).
-     *
-     * @param {string} agentName - The name of the agent to add lazily.
-     */
-    addLazyAgent(agentName: string): void;
-    /**
-     * Sets up agents in the crew by name.
-     * For an array of Agent names provided, it adds
-     * the agent to the crew if not already present.
-     * @param {string[]} agentNames - An array of agent names to configure.
-     * @returns {Map<string, StatefulAxAgent> | null} A map of agent names to their corresponding instances.
-     */
-    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.
-     */
-    destroy(): void;
-    /**
-     * Resets all cost and usage tracking for the entire crew.
-     * Also calls each agent's `resetUsage` (if available) and clears crew-level metrics.
-     */
-    resetCosts(): void;
-    /**
-     * Get an aggregate metrics snapshot for the entire crew.
-     * Sums requests, errors, tokens, and estimated cost across all agents in the crew.
-     *
-     * @returns Crew-level metrics snapshot.
-     */
-    getCrewMetrics(): import("../metrics/types.js").MetricsSnapshot;
-    /**
-     * Reset all tracked metrics for the entire crew.
-     * Use to clear totals before a new measurement period.
-     */
-    resetCrewMetrics(): void;
-}
-export { AxCrew, LazyStatefulAxAgent };
-export type { StatefulAxAgent };
+export { StatefulAxAgent } from './statefulAgent.js';
+export type { ParsedAgentConfig } from './statefulAgent.js';
+export { LazyStatefulAxAgent } from './lazyAgent.js';
+export { AxCrew } from './crew.js';