@amitdeshmukh/ax-crew 3.11.3 → 4.0.1

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.
@@ -155,7 +155,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 +167,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();

package/examples/mcp-agent.ts

@@ -1,10 +1,11 @@
 import { AxCrew } from "../dist/index.js";
+import type { AxCrewConfig } from "../dist/index.js";
 
 import dotenv from "dotenv";
 dotenv.config();
 
 // Define the crew configuration
-const config = {
+const config: AxCrewConfig = {
   crew: [
     {
       name: "MapsAgent",
@@ -15,9 +16,11 @@ const config = {
       ai: {
         model: "claude-3-5-sonnet-latest",
         temperature: 0,
+        maxTokens: 1000,
+        stream: true
       },
       options: {
-        debug: true,
+        debug: true
       },
       "mcpServers": {
         "google-maps": {
@@ -43,6 +46,7 @@ const config = {
         model: "gpt-4o-mini",
         maxTokens: 1000,
         temperature: 0,
+        stream: true
       },
       options: {
         debug: true,
@@ -59,10 +63,10 @@ const config = {
       ai: {
         model: "gemini-1.5-pro",
         temperature: 0,
+        stream: true
       },
       options: {
         debug: false,
-        codeExecution: true,
       },
     },
   ],
@@ -89,11 +93,11 @@ const main = async (): Promise<void> => {
 
     console.log(`\nAnswer: ${JSON.stringify(managerResponse?.answer, null, 2)}`);
 
-    // Print usage costs
-    console.log("\nUsage:\n+++++++++++++++++++++++++++++++++");
-    console.log("Manager Agent Cost in $:", JSON.stringify(managerAgent?.getAccumulatedCosts()?.totalCost, null, 2));
-    console.log("Maps Agent Cost in $:", JSON.stringify(mapsAgent?.getAccumulatedCosts()?.totalCost, null, 2));
-    console.log("Total Cost in $:", JSON.stringify(crew.getAggregatedCosts()?.totalCost, null, 2));
+    // Print metrics
+    console.log("\nMetrics:\n+++++++++++++++++++++++++++++++++");
+    console.log("Manager Agent Metrics:", JSON.stringify((managerAgent as any)?.getMetrics?.(), null, 2));
+    console.log("Maps Agent Metrics:", JSON.stringify((mapsAgent as any)?.getMetrics?.(), null, 2));
+    console.log("Crew Metrics:", JSON.stringify((crew as any)?.getCrewMetrics?.(), null, 2));
 };
 
 main()

package/examples/solve-math-problem.ts

@@ -1,7 +1,8 @@
 import { AxCrew } from "../dist/index.js";
+import type { AxCrewConfig } from "../src/index.js";
 
 // Define the crew configuration
-const config = {
+const config: AxCrewConfig = {
   crew: [
     {
       name: "MathAgent",
@@ -62,10 +63,11 @@ const main = async (): Promise<void> => {
 
     console.log(`\nAnswer: ${JSON.stringify(managerResponse.answer, null, 2)}`);
 
-    // Print usage costs
-    console.log("\nUsage:\n+++++++++++++++++++++++++++++++++");
-    console.log("Manager Agent:", JSON.stringify(managerAgent.getAccumulatedCosts(), null, 2));
-    console.log("Total Cost:", JSON.stringify(crew.getAggregatedCosts(), null, 2));
+    // Print metrics
+    console.log("\nMetrics:\n+++++++++++++++++++++++++++++++++");
+    console.log("Manager Agent Metrics:", JSON.stringify((managerAgent as any).getMetrics?.(), null, 2));
+    console.log("Math Agent Metrics:", JSON.stringify((mathAgent as any).getMetrics?.(), null, 2));
+    console.log("Crew Metrics:", JSON.stringify((crew as any).getCrewMetrics?.(), null, 2));
   }
 };
 

package/examples/streaming.ts

@@ -1,10 +1,11 @@
 import { AxCrew } from "../dist/index.js";
+import type { AxCrewConfig } from "../src/index.js";
 
 import dotenv from "dotenv";
 dotenv.config();
 
 // Define the crew configuration
-const config = {
+const config: AxCrewConfig = {
   crew: [
     {
       name: "ManagerAgent",
@@ -75,11 +76,11 @@ const main = async (): Promise<void> => {
       }
     }
 
-    // Print usage costs
-    console.log("\nUsage:\n+++++++++++++++++++++++++++++++++");
-    console.log("Manager Agent Cost in $:", JSON.stringify(managerAgent?.getAccumulatedCosts()?.totalCost, null, 2));
-    console.log("Math Agent Cost in $:", JSON.stringify(mathAgent?.getAccumulatedCosts()?.totalCost, null, 2));
-    console.log("Total Cost in $:", JSON.stringify(crew.getAggregatedCosts()?.totalCost, null, 2));
+    // Print metrics
+    console.log("\nMetrics:\n+++++++++++++++++++++++++++++++++");
+    console.log("Manager Agent Metrics:", JSON.stringify((managerAgent as any)?.getMetrics?.(), null, 2));
+    console.log("Math Agent Metrics:", JSON.stringify((mathAgent as any)?.getMetrics?.(), null, 2));
+    console.log("Crew Metrics:", JSON.stringify((crew as any)?.getCrewMetrics?.(), null, 2));
 };
 
 main()

package/package.json

@@ -1,7 +1,7 @@
 {
   "type": "module",
   "name": "@amitdeshmukh/ax-crew",
-  "version": "3.11.3",
+  "version": "4.0.1",
   "description": "Build and launch a crew of AI agents with shared state. Built with axllm.dev",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
@@ -17,17 +17,19 @@
     "test:ui": "vitest --ui"
   },
   "dependencies": {
+    "big.js": "^7.0.1",
     "decimal.js": "^10.5.0",
     "dotenv": "^16.4.5",
     "upgrade": "^1.1.0",
     "uuid": "^10.0.0"
   },
   "peerDependencies": {
-    "@ax-llm/ax": "^13.0.8",
-    "@ax-llm/ax-tools": "^13.0.8"
+    "@ax-llm/ax": "14.0.16",
+    "@ax-llm/ax-tools": "14.0.16"
   },
   "devDependencies": {
     "@testing-library/jest-dom": "^6.6.3",
+    "@types/big.js": "^6.2.2",
     "@types/node": "^20.14.9",
     "@types/uuid": "^10.0.0",
     "@vitest/coverage-v8": "^3.0.9",

package/src/agents/agentConfig.ts

@@ -1,10 +1,8 @@
 import fs from 'fs';
-// Import all the providers
-import { AxAIAnthropic, AxAIOpenAI, AxAIAzureOpenAI, AxAICohere, AxAIDeepSeek, AxAIGoogleGemini, AxAIGroq, AxAIHuggingFace, AxAIMistral, AxAIOllama, AxAITogether, AxAIReka, AxAIGrok } from '@ax-llm/ax';
-// Import Ax types
+// Import Ax factory and MCP transports (as exported by current package)
+import { ai, AxMCPClient, AxMCPHTTPSSETransport, AxMCPStreambleHTTPTransport, AxDefaultCostTracker } from '@ax-llm/ax'
 import type { AxFunction } from '@ax-llm/ax';
-// Import the MCP client and transports
-import { AxMCPClient, AxMCPHTTPSSETransport, AxMCPStreambleHTTPTransport } from '@ax-llm/ax'
+// STDIO transport from tools package
 import { AxMCPStdioTransport } from '@ax-llm/ax-tools'
 import { PROVIDER_API_KEYS } from '../config/index.js';
 import type { 
@@ -14,28 +12,29 @@ import type {
   MCPTransportConfig, 
   MCPStdioTransportConfig, 
   MCPHTTPSSETransportConfig,
-  MCPStreambleHTTPTransportConfig
+  MCPStreamableHTTPTransportConfig
 } from '../types.js';
+import type { Provider } from '../types.js';
 
-// Define a mapping from provider names to their respective constructors
-const AIConstructors: Record<string, any> = {
-  'anthropic': AxAIAnthropic,
-  'azure-openai': AxAIAzureOpenAI,
-  'cohere': AxAICohere,
-  'deepseek': AxAIDeepSeek,
-  'google-gemini': AxAIGoogleGemini,
-  'groq': AxAIGroq,
-  'huggingFace': AxAIHuggingFace,
-  'mistral': AxAIMistral,
-  'ollama': AxAIOllama,
-  'openai': AxAIOpenAI,
-  'together': AxAITogether,
-  'reka': AxAIReka,
-  'grok': AxAIGrok
-};
+// Canonical provider slugs supported by ai() factory
+const PROVIDER_CANONICAL = new Set([
+  'openai',
+  'anthropic',
+  'google-gemini',
+  'mistral',
+  'groq',
+  'cohere',
+  'together',
+  'deepseek',
+  'ollama',
+  'huggingface',
+  'openrouter',
+  'azure-openai',
+  'reka',
+  'x-grok'
+]);
 
-// Provider type
-export type Provider = keyof typeof AIConstructors;
+// Provider type lives in src/types.ts
 
 // Type guard to check if config is stdio transport
 export function isStdioTransport(config: MCPTransportConfig): config is MCPStdioTransportConfig {
@@ -48,7 +47,7 @@ export function isHTTPSSETransport(config: MCPTransportConfig): config is MCPHTT
 }
 
 // Type guard to check if config is streamable HTTP transport
-export function isStreambleHTTPTransport(config: MCPTransportConfig): config is MCPStreambleHTTPTransportConfig {
+export function isStreambleHTTPTransport(config: MCPTransportConfig): config is MCPStreamableHTTPTransportConfig {
   return 'mcpEndpoint' in config;
 }
 
@@ -205,11 +204,12 @@ const parseAgentConfig = async (
       throw new Error(`AI agent with name ${agentName} is not configured`);
     }
 
-    // Get the constructor for the AI agent's provider
-    const AIConstructor = AIConstructors[agentConfigData.provider];
-    if (!AIConstructor) {
-      throw new Error(`AI provider ${agentConfigData.provider} is not supported. Did you mean '${agentConfigData.provider.toLowerCase()}'?`);
+    // Enforce canonical provider slug
+    const lower = agentConfigData.provider.toLowerCase();
+    if (!PROVIDER_CANONICAL.has(lower)) {
+      throw new Error(`AI provider ${agentConfigData.provider} is not supported. Use one of: ${Array.from(PROVIDER_CANONICAL).join(', ')}`);
     }
+    const provider = lower as Provider;
 
     // If an API Key property is present, get the API key for the AI agent from the environment variables
     let apiKey = '';
@@ -223,25 +223,30 @@ const parseAgentConfig = async (
       throw new Error(`Provider key name is missing in the agent configuration`);
     }
 
-    // Create an instance of the AI agent and set options
-    const ai = new AIConstructor({
+    // Create a cost tracker instance and pass to ai()
+    const costTracker = new AxDefaultCostTracker();
+
+    // Create an instance of the AI agent via factory
+    const aiArgs: any = {
+      name: provider,
       apiKey,
       config: agentConfigData.ai,
       options: {
         debug: agentConfigData.debug || false,
-        ...agentConfigData.options
+        ...agentConfigData.options,
+        // Attach default cost tracker so usage/costs are recorded by provider layer
+        trackers: [costTracker]
       }
-    });
-    // If an apiURL is provided in the agent config, set it in the AI agent
+    };
     if (agentConfigData.apiURL) {
       try {
-        // Validate apiURL format
         new URL(agentConfigData.apiURL);
-        ai.setAPIURL(agentConfigData.apiURL);
+        aiArgs.apiURL = agentConfigData.apiURL;
       } catch (error) {
         throw new Error(`Invalid apiURL provided: ${agentConfigData.apiURL}`);
-      }      
+      }
     }
+    const aiInstance = ai(aiArgs);
 
     // If an mcpServers config is provided in the agent config, convert to functions
     const mcpFunctions = await initializeMCPServers(agentConfigData);
@@ -265,13 +270,14 @@ const parseAgentConfig = async (
     
     // Return AI instance and Agent parameters
     return {
-      ai,
+      ai: aiInstance,
       name: agentName,
       description: agentConfigData.description,
       signature: agentConfigData.signature,
       functions: agentFunctions,
       subAgentNames: agentConfigData.agents || [],
       examples: agentConfigData.examples || [],
+      tracker: costTracker,
     };
   } catch (error) {
     if (error instanceof Error) {