@amitdeshmukh/ax-crew 3.11.3 → 4.1.1

package/dist/agents/index.js

@@ -2,12 +2,20 @@ import { v4 as uuidv4 } from "uuid";
 import { AxAgent } from "@ax-llm/ax";
 import { createState } from "../state/index.js";
 import { parseCrewConfig, parseAgentConfig } from "./agentConfig.js";
-import { StateFulAxAgentUsage } from "./agentUseCosts.js";
+import { MetricsRegistry } from "../metrics/index.js";
 // Extend the AxAgent class from ax-llm
 class StatefulAxAgent extends AxAgent {
     state;
     axai;
     agentName;
+    costTracker;
+    lastRecordedCostUSD = 0;
+    isAxAIService(obj) {
+        return !!obj && typeof obj.getName === 'function' && typeof obj.chat === 'function';
+    }
+    isAxAIInstance(obj) {
+        return !!obj && typeof obj === 'object' && ('defaults' in obj || 'modelInfo' in obj);
+    }
     constructor(ai, options, state) {
         const { examples, ...restOptions } = options;
         const formattedOptions = {
@@ -26,9 +34,12 @@ class StatefulAxAgent extends AxAgent {
     // Implementation
     async forward(first, second, third) {
         let result;
+        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
-        if ('apiURL' in first) {
+        if (this.isAxAIService(first)) {
             // Sub-agent case (called with AI service)
             result = await super.forward(this.axai, second, third);
         }
@@ -36,17 +47,51 @@ class StatefulAxAgent extends AxAgent {
             // Direct call case
             result = await super.forward(this.axai, first, second);
         }
-        // Track costs after the call
-        const cost = this.getLastUsageCost();
-        if (cost) {
-            StateFulAxAgentUsage.trackCostInState(this.agentName, cost, this.state);
+        // Track metrics and costs after the call using built-in usage
+        const durationMs = performance.now() - start;
+        MetricsRegistry.recordRequest(labels, false, durationMs);
+        // Always record tokens from built-in usage array if present
+        const builtIn = this.getUsage?.();
+        if (Array.isArray(builtIn)) {
+            const totals = builtIn.reduce((acc, u) => {
+                const pt = u.tokens?.promptTokens ?? u.promptTokens ?? 0;
+                const ct = u.tokens?.completionTokens ?? u.completionTokens ?? 0;
+                acc.promptTokens += typeof pt === 'number' ? pt : 0;
+                acc.completionTokens += typeof ct === 'number' ? ct : 0;
+                // also aggregate per-model to feed Ax tracker
+                const model = u.model || this.axai?.getLastUsedChatModel?.() || this.axai?.defaults?.model;
+                if (model) {
+                    acc.byModel[model] = (acc.byModel[model] || 0) + (pt + ct);
+                }
+                return acc;
+            }, { promptTokens: 0, completionTokens: 0, byModel: {} });
+            MetricsRegistry.recordTokens(labels, {
+                promptTokens: totals.promptTokens,
+                completionTokens: totals.completionTokens,
+                totalTokens: totals.promptTokens + totals.completionTokens,
+            });
+            // Feed Ax's cost tracker with token totals per model; Ax owns pricing
+            const costTracker = this.costTracker;
+            try {
+                for (const [m, count] of Object.entries(totals.byModel)) {
+                    costTracker?.trackTokens?.(count, m);
+                }
+                const totalUSD = Number(costTracker?.getCurrentCost?.() ?? 0);
+                if (!Number.isNaN(totalUSD) && totalUSD > 0) {
+                    MetricsRegistry.recordEstimatedCost(labels, totalUSD);
+                }
+            }
+            catch { }
         }
         return result;
     }
     // Implementation
     streamingForward(first, second, third) {
+        const start = performance.now();
+        const crewId = this.state?.crewId || (this.state.get?.('crewId')) || 'default';
+        const labels = { crewId, agent: this.agentName };
         let streamingResult;
-        if ('apiURL' in first) {
+        if (this.isAxAIService(first)) {
             streamingResult = super.streamingForward(this.axai, second, third);
         }
         else {
@@ -60,34 +105,64 @@ class StatefulAxAgent extends AxAgent {
                 }
             }
             finally {
-                const cost = this.getLastUsageCost();
-                if (cost) {
-                    StateFulAxAgentUsage.trackCostInState(this.agentName, cost, this.state);
+                const durationMs = performance.now() - start;
+                MetricsRegistry.recordRequest(labels, true, durationMs);
+                // Record tokens from built-in usage array if present
+                const builtIn = this.getUsage?.();
+                if (Array.isArray(builtIn)) {
+                    const totals = builtIn.reduce((acc, u) => {
+                        const pt = u.tokens?.promptTokens ?? u.promptTokens ?? 0;
+                        const ct = u.tokens?.completionTokens ?? u.completionTokens ?? 0;
+                        acc.promptTokens += typeof pt === 'number' ? pt : 0;
+                        acc.completionTokens += typeof ct === 'number' ? ct : 0;
+                        const model = u.model || this.axai?.getLastUsedChatModel?.() || this.axai?.defaults?.model;
+                        if (model) {
+                            acc.byModel[model] = (acc.byModel[model] || 0) + (pt + ct);
+                        }
+                        return acc;
+                    }, { promptTokens: 0, completionTokens: 0, byModel: {} });
+                    MetricsRegistry.recordTokens(labels, {
+                        promptTokens: totals.promptTokens,
+                        completionTokens: totals.completionTokens,
+                        totalTokens: totals.promptTokens + totals.completionTokens,
+                    });
+                    const costTracker = this.costTracker;
+                    try {
+                        for (const [m, count] of Object.entries(totals.byModel)) {
+                            costTracker?.trackTokens?.(count, m);
+                        }
+                        const totalUSD = Number(costTracker?.getCurrentCost?.() ?? 0);
+                        if (!Number.isNaN(totalUSD) && totalUSD > 0) {
+                            MetricsRegistry.recordEstimatedCost(labels, totalUSD);
+                        }
+                    }
+                    catch { }
+                }
+                // Record estimated cost (USD) via attached tracker if available
+                const costTracker = this.costTracker;
+                try {
+                    const totalUSD = Number(costTracker?.getCurrentCost?.() ?? 0);
+                    if (!Number.isNaN(totalUSD) && totalUSD > 0) {
+                        MetricsRegistry.recordEstimatedCost(labels, totalUSD);
+                    }
                 }
+                catch { }
             }
         }).bind(this)();
         return wrappedGenerator;
     }
-    // Get the usage cost for the most recent run of the agent
-    getLastUsageCost() {
-        const { modelUsage, modelInfo, defaults } = this.axai;
-        // Check if all required properties exist
-        if (!modelUsage?.promptTokens || !modelUsage?.completionTokens) {
-            return null;
-        }
-        if (!modelInfo || !defaults?.model) {
-            return null;
-        }
-        const currentModelInfo = modelInfo.find((m) => m.name === defaults.model);
-        if (!currentModelInfo?.promptTokenCostPer1M || !currentModelInfo?.completionTokenCostPer1M) {
-            return null;
-        }
-        return StateFulAxAgentUsage.calculateCost(modelUsage, currentModelInfo);
-    }
+    // Legacy cost API removed: rely on Ax trackers for cost reporting
+    getLastUsageCost() { return null; }
     // Get the accumulated costs for all runs of this agent
-    getAccumulatedCosts() {
-        const stateKey = `${StateFulAxAgentUsage.STATE_KEY_PREFIX}${this.agentName}`;
-        return this.state.get(stateKey);
+    getAccumulatedCosts() { return null; }
+    // Metrics API for this agent
+    getMetrics() {
+        const crewId = this.state?.crewId || (this.state.get?.('crewId')) || 'default';
+        return MetricsRegistry.snapshot({ crewId, agent: this.agentName });
+    }
+    resetMetrics() {
+        const crewId = this.state?.crewId || (this.state.get?.('crewId')) || 'default';
+        MetricsRegistry.reset({ crewId, agent: this.agentName });
     }
 }
 /**
@@ -121,6 +196,8 @@ class AxCrew {
         this.crewId = crewId;
         this.agents = new Map();
         this.state = createState(crewId);
+        // Make crewId discoverable to metrics
+        this.state.set('crewId', crewId);
     }
     /**
      * Factory function for creating an agent.
@@ -132,7 +209,7 @@ class AxCrew {
         try {
             const agentConfig = await parseAgentConfig(agentName, this.crewConfig, this.functionsRegistry, this.state);
             // Destructure with type assertion
-            const { ai, name, description, signature, functions, subAgentNames, examples } = agentConfig;
+            const { ai, name, description, signature, functions, subAgentNames, examples, tracker } = agentConfig;
             // Get subagents for the AI agent
             const subAgents = subAgentNames.map((subAgentName) => {
                 if (!this.agents?.get(subAgentName)) {
@@ -140,15 +217,40 @@ class AxCrew {
                 }
                 return this.agents?.get(subAgentName);
             });
+            // Dedupe sub-agents by name (defensive)
+            const subAgentSet = new Map();
+            for (const sa of subAgents.filter((agent) => agent !== undefined)) {
+                const n = sa?.agentName ?? sa?.name ?? '';
+                if (!subAgentSet.has(n))
+                    subAgentSet.set(n, sa);
+            }
+            const uniqueSubAgents = Array.from(subAgentSet.values());
+            // Dedupe functions by name and avoid collision with sub-agent names
+            const subAgentNameSet = new Set(uniqueSubAgents.map((sa) => sa?.agentName ?? sa?.name).filter(Boolean));
+            const uniqueFunctions = [];
+            const seenFn = new Set();
+            for (const fn of functions.filter((fn) => fn !== undefined)) {
+                const fnName = fn.name;
+                if (subAgentNameSet.has(fnName)) {
+                    // Skip function that collides with a sub-agent name
+                    continue;
+                }
+                if (!seenFn.has(fnName)) {
+                    seenFn.add(fnName);
+                    uniqueFunctions.push(fn);
+                }
+            }
             // Create an instance of StatefulAxAgent
             const agent = new StatefulAxAgent(ai, {
                 name,
                 description,
+                definition: agentConfig.definition,
                 signature,
-                functions: functions.filter((fn) => fn !== undefined),
-                agents: subAgents.filter((agent) => agent !== undefined),
+                functions: uniqueFunctions,
+                agents: uniqueSubAgents,
                 examples,
             }, this.state);
+            agent.costTracker = tracker;
             return agent;
         }
         catch (error) {
@@ -273,18 +375,31 @@ class AxCrew {
         this.agents = null;
         this.state.reset();
     }
-    /**
-     * Gets aggregated costs for all agents in the crew
-     * @returns Aggregated cost information for all agents
-     */
-    getAggregatedCosts() {
-        return StateFulAxAgentUsage.getAggregatedCosts(this.state);
-    }
     /**
      * Resets all cost tracking for the crew
      */
     resetCosts() {
-        StateFulAxAgentUsage.resetCosts(this.state);
+        // Reset AxAgent built-in usage and our metrics registry
+        if (this.agents) {
+            for (const [, agent] of this.agents) {
+                try {
+                    agent.resetUsage?.();
+                }
+                catch { }
+                try {
+                    agent.resetMetrics?.();
+                }
+                catch { }
+            }
+        }
+        MetricsRegistry.reset({ crewId: this.crewId });
+    }
+    // Metrics API
+    getCrewMetrics() {
+        return MetricsRegistry.snapshotCrew(this.crewId);
+    }
+    resetCrewMetrics() {
+        MetricsRegistry.reset({ crewId: this.crewId });
     }
 }
 export { AxCrew };

package/dist/index.d.ts

@@ -2,6 +2,8 @@ import { AxCrew } from './agents/index.js';
 import { AxCrewFunctions } from './functions/index.js';
 import type { CrewConfigInput, AgentConfig } from './types.js';
 import type { UsageCost, AggregatedMetrics, AggregatedCosts, StateInstance, FunctionRegistryType } from './types.js';
+export * from './metrics/index.js';
+export { MetricsRegistry } from './metrics/index.js';
 /**
  * The configuration for an AxCrew.
  *

package/dist/index.js

@@ -1,3 +1,5 @@
 import { AxCrew } from './agents/index.js';
 import { AxCrewFunctions } from './functions/index.js';
+export * from './metrics/index.js';
+export { MetricsRegistry } from './metrics/index.js';
 export { AxCrew, AxCrewFunctions, };

package/dist/metrics/index.d.ts

@@ -0,0 +1,2 @@
+export * from './types.js';
+export * as MetricsRegistry from './registry.js';

package/dist/metrics/index.js

@@ -0,0 +1,2 @@
+export * from './types.js';
+export * as MetricsRegistry from './registry.js';

package/dist/metrics/registry.d.ts

@@ -0,0 +1,9 @@
+import type { LabelKeys, MetricsSnapshot, TokenUsage } from './types.js';
+export declare function recordRequest(labels: LabelKeys, streaming: boolean, durationMs: number): void;
+export declare function recordError(labels: LabelKeys): void;
+export declare function recordTokens(labels: LabelKeys, usage: TokenUsage): void;
+export declare function recordEstimatedCost(labels: LabelKeys, usd: number): void;
+export declare function recordFunctionCall(labels: LabelKeys, latencyMs: number): void;
+export declare function snapshot(labels: LabelKeys): MetricsSnapshot;
+export declare function reset(labels?: LabelKeys): void;
+export declare function snapshotCrew(crewId: string): MetricsSnapshot;

package/dist/metrics/registry.js

@@ -0,0 +1,138 @@
+import Big from 'big.js';
+const store = new Map();
+function keyOf(labels) {
+    const { crewId, agent = '', provider = '', model = '' } = labels;
+    return [crewId, agent, provider, model].join('|');
+}
+function getOrInit(labels) {
+    const k = keyOf(labels);
+    let c = store.get(k);
+    if (!c) {
+        c = {
+            requests: 0,
+            errors: 0,
+            streaming: 0,
+            durationMsSum: 0,
+            durationCount: 0,
+            inputTokens: 0,
+            outputTokens: 0,
+            estimatedCostUSD: 0,
+            functionCalls: 0,
+            functionLatencyMs: 0,
+        };
+        store.set(k, c);
+    }
+    return c;
+}
+export function recordRequest(labels, streaming, durationMs) {
+    const c = getOrInit(labels);
+    c.requests += 1;
+    if (streaming)
+        c.streaming += 1;
+    c.durationMsSum += durationMs;
+    c.durationCount += 1;
+}
+export function recordError(labels) {
+    const c = getOrInit(labels);
+    c.errors += 1;
+}
+export function recordTokens(labels, usage) {
+    const c = getOrInit(labels);
+    c.inputTokens += usage.promptTokens || 0;
+    c.outputTokens += usage.completionTokens || 0;
+}
+export function recordEstimatedCost(labels, usd) {
+    const c = getOrInit(labels);
+    const current = new Big(c.estimatedCostUSD || 0);
+    const addition = new Big(usd || 0);
+    c.estimatedCostUSD = Number(current.plus(addition));
+}
+export function recordFunctionCall(labels, latencyMs) {
+    const c = getOrInit(labels);
+    c.functionCalls += 1;
+    c.functionLatencyMs += latencyMs || 0;
+}
+export function snapshot(labels) {
+    const c = getOrInit(labels);
+    const totalTokens = c.inputTokens + c.outputTokens;
+    return {
+        provider: labels.provider,
+        model: labels.model,
+        requests: {
+            totalRequests: c.requests,
+            totalErrors: c.errors,
+            errorRate: c.requests > 0 ? c.errors / c.requests : 0,
+            totalStreamingRequests: c.streaming,
+            durationMsSum: c.durationMsSum,
+            durationCount: c.durationCount,
+        },
+        tokens: {
+            promptTokens: c.inputTokens,
+            completionTokens: c.outputTokens,
+            totalTokens,
+        },
+        estimatedCostUSD: Number(new Big(c.estimatedCostUSD || 0).round(5)),
+        functions: {
+            totalFunctionCalls: c.functionCalls,
+            totalFunctionLatencyMs: c.functionLatencyMs,
+        },
+    };
+}
+export function reset(labels) {
+    if (!labels) {
+        store.clear();
+        return;
+    }
+    const k = keyOf(labels);
+    store.delete(k);
+}
+export function snapshotCrew(crewId) {
+    const empty = {
+        requests: 0,
+        errors: 0,
+        streaming: 0,
+        durationMsSum: 0,
+        durationCount: 0,
+        inputTokens: 0,
+        outputTokens: 0,
+        estimatedCostUSD: 0,
+        functionCalls: 0,
+        functionLatencyMs: 0,
+    };
+    const agg = Array.from(store.entries()).reduce((acc, [k, v]) => {
+        if (k.startsWith(crewId + '|')) {
+            acc.requests += v.requests;
+            acc.errors += v.errors;
+            acc.streaming += v.streaming;
+            acc.durationMsSum += v.durationMsSum;
+            acc.durationCount += v.durationCount;
+            acc.inputTokens += v.inputTokens;
+            acc.outputTokens += v.outputTokens;
+            acc.estimatedCostUSD = Number(new Big(acc.estimatedCostUSD || 0).plus(v.estimatedCostUSD || 0));
+            acc.functionCalls += v.functionCalls;
+            acc.functionLatencyMs += v.functionLatencyMs;
+        }
+        return acc;
+    }, { ...empty });
+    const totalTokens = agg.inputTokens + agg.outputTokens;
+    return {
+        requests: {
+            totalRequests: agg.requests,
+            totalErrors: agg.errors,
+            errorRate: agg.requests > 0 ? agg.errors / agg.requests : 0,
+            totalStreamingRequests: agg.streaming,
+            durationMsSum: agg.durationMsSum,
+            durationCount: agg.durationCount,
+        },
+        tokens: {
+            promptTokens: agg.inputTokens,
+            completionTokens: agg.outputTokens,
+            totalTokens,
+        },
+        estimatedCostUSD: Number(new Big(agg.estimatedCostUSD || 0).round(5)),
+        functions: {
+            totalFunctionCalls: agg.functionCalls,
+            totalFunctionLatencyMs: agg.functionLatencyMs,
+        },
+    };
+}

package/dist/metrics/types.d.ts

@@ -0,0 +1,40 @@
+export interface TokenUsage {
+    promptTokens: number;
+    completionTokens: number;
+    totalTokens?: number;
+}
+export interface CostSnapshot {
+    usdTotal: number;
+    tokenUsage: TokenUsage;
+}
+export interface RequestStats {
+    totalRequests: number;
+    totalErrors: number;
+    errorRate: number;
+    totalStreamingRequests: number;
+    durationMsSum: number;
+    durationCount: number;
+}
+export interface FunctionStats {
+    totalFunctionCalls: number;
+    totalFunctionLatencyMs: number;
+}
+export interface MetricsSnapshot {
+    provider?: string;
+    model?: string;
+    requests: RequestStats;
+    tokens: TokenUsage;
+    estimatedCostUSD: number;
+    functions: FunctionStats;
+}
+export interface LabelKeys {
+    crewId: string;
+    agent?: string;
+    provider?: string;
+    model?: string;
+}
+export interface BudgetConfig {
+    maxTokens?: number;
+    maxCost?: number;
+    costPerModel?: Record<string, number>;
+}

package/dist/metrics/types.js

@@ -0,0 +1 @@
+export {};

package/dist/types.d.ts

@@ -1,5 +1,5 @@
 import type { AxFunction, AxSignature, AxModelConfig, AxMCPStreamableHTTPTransportOptions, AxProgramForwardOptions } from '@ax-llm/ax';
-import type { Provider } from './agents/agentConfig.js';
+export type Provider = 'openai' | 'anthropic' | 'google-gemini' | 'mistral' | 'groq' | 'cohere' | 'together' | 'deepseek' | 'ollama' | 'huggingface' | 'openrouter' | 'azure-openai' | 'reka' | 'x-grok';
 /**
  * A state instance that is shared between agents.
  * This can be used to store data that becomes available to all agents and functions in an out-of-band manner.
@@ -117,7 +117,7 @@ interface MCPHTTPSSETransportConfig {
  * @property {string} mcpEndpoint - The HTTP endpoint URL for the MCP server.
  * @property {AxMCPStreamableHTTPTransportOptions} options - Optional transport options.
  */
-interface MCPStreambleHTTPTransportConfig {
+interface MCPStreamableHTTPTransportConfig {
     mcpEndpoint: string;
     options?: AxMCPStreamableHTTPTransportOptions;
 }
@@ -126,7 +126,7 @@ interface MCPStreambleHTTPTransportConfig {
  *
  * @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 | MCPStreambleHTTPTransportConfig;
+type MCPTransportConfig = MCPStdioTransportConfig | MCPHTTPSSETransportConfig | MCPStreamableHTTPTransportConfig;
 /**
  * The configuration for an agent.
  *
@@ -138,7 +138,7 @@ type MCPTransportConfig = MCPStdioTransportConfig | MCPHTTPSSETransportConfig |
  * @property {AxModelConfig & { model: string }} ai - The AI model configuration to be passed to the agent.
  * @property {boolean} debug - Whether to enable debug mode.
  * @property {string} apiURL - Set this if you are using a custom API URL e.g. ollama on localhost.
- * @property {Partial<AxProgramForwardOptions>} options - Agent options including thinkingTokenBudget, showThoughts, etc. Refer to the Ax documentation for more details.
+ * @property {Partial<AxProgramForwardOptions<any>> & Record<string, any>} options - Agent options including thinkingTokenBudget, showThoughts, etc. Also allows arbitrary provider-specific keys.
  * @property {string[]} functions - Function names to be used by the agent.
  * @property {string[]} agents - Sub-agent available to the agent.
  * @property {Record<string, any>[]} examples - DSPy examples for the agent to learn from.
@@ -147,6 +147,16 @@ type MCPTransportConfig = MCPStdioTransportConfig | MCPHTTPSSETransportConfig |
 interface AgentConfig {
     name: string;
     description: string;
+    /**
+     * Optional detailed persona/program definition. If provided, becomes the system prompt.
+     * Must be at least 100 characters per Ax semantics.
+     */
+    definition?: string;
+    /**
+     * Optional alias for definition for clarity. If provided and definition is omitted,
+     * this will be used as the program definition/system prompt.
+     */
+    prompt?: string;
     signature: string | AxSignature;
     provider: Provider;
     providerKeyName?: string;
@@ -155,7 +165,7 @@ interface AgentConfig {
     };
     debug?: boolean;
     apiURL?: string;
-    options?: Partial<AxProgramForwardOptions>;
+    options?: Partial<AxProgramForwardOptions<any>> & Record<string, any>;
     functions?: string[];
     agents?: string[];
     examples?: Array<Record<string, any>>;
@@ -167,4 +177,4 @@ interface AgentConfig {
 type CrewConfigInput = string | {
     crew: AgentConfig[];
 };
-export { type AgentConfig, type CrewConfigInput, type AggregatedMetrics, type StateInstance, type FunctionRegistryType, type MCPStdioTransportConfig, type MCPHTTPSSETransportConfig, type MCPStreambleHTTPTransportConfig, type MCPTransportConfig, type ModelUsage, type ModelInfo, type UsageCost, type AggregatedCosts };
+export { type AgentConfig, type CrewConfigInput, type AggregatedMetrics, type StateInstance, type FunctionRegistryType, type MCPStdioTransportConfig, type MCPHTTPSSETransportConfig, type MCPStreamableHTTPTransportConfig, type MCPTransportConfig, type ModelUsage, type ModelInfo, type UsageCost, type AggregatedCosts };

package/examples/basic-researcher-writer.ts

@@ -1,17 +1,21 @@
 import { AxCrew } from "../dist/index.js";
-import { AxCrewFunctions } from "../src/functions/index.js";
+import { AxCrewFunctions } from "../dist/functions/index.js";
+import type { AxCrewConfig } from "../dist/index.js";
+import type { Provider } from "../dist/types.js";
 
 // Example agent configuration
-const agentConfig = {
+const agentConfig: AxCrewConfig = {
   crew: [
     {
       name: "researcher",
       description: "A research agent that finds information",
       signature: "query:string -> research:string",
-      provider: "anthropic",
-      providerKeyName: "ANTHROPIC_API_KEY",
+      provider: "google-gemini" as Provider,
+      providerKeyName: "GEMINI_API_KEY",
       ai: {
-        model: "claude-3-haiku-20240307"
+        model: "gemini-2.5-flash-lite",
+        maxTokens: 4000,
+        stream: true
       },
       options: {
         debug: true,
@@ -22,10 +26,12 @@ const agentConfig = {
       name: "writer",
       description: "A writing agent that creates content",
       signature: "topic:string -> article:string",
-      provider: "anthropic",
+      provider: "anthropic" as Provider,
       providerKeyName: "ANTHROPIC_API_KEY",
       ai: {
-        model: "claude-3-haiku-20240307"
+        model: "claude-3-haiku-20240307",
+        maxTokens: 4000,
+        stream: true
       },
       options: {
         debug: true,
@@ -61,12 +67,11 @@ async function main() {
     // Print the article
     console.log("Article:", article);
     
-    // Print usage costs
-    console.log("\nUsage:\n+++++++++++++++++++++++++++++++++");
-    console.log("Writer Agent:", JSON.stringify(writer.getAccumulatedCosts(), null, 2));
-    console.log("Researcher Agent Last Usage:", JSON.stringify(researcher.getLastUsageCost(), null, 2));
-    console.log("Researcher Agent Accumulated:", JSON.stringify(researcher.getAccumulatedCosts(), null, 2));
-    console.log("Total Cost:", JSON.stringify(crew.getAggregatedCosts(), null, 2));
+    // Print metrics snapshots (new mechanism)
+    console.log("\nMetrics:\n+++++++++++++++++++++++++++++++++");
+    console.log("Writer Metrics:", JSON.stringify((writer as any).getMetrics?.(), null, 2));
+    console.log("Researcher Metrics:", JSON.stringify((researcher as any).getMetrics?.(), null, 2));
+    console.log("Crew Metrics:", JSON.stringify((crew as any).getCrewMetrics?.(), null, 2));
 
     // If you want to start fresh with cost tracking
     crew.resetCosts();