@amitdeshmukh/ax-crew 4.1.2 → 6.0.0

package/package.json

@@ -1,7 +1,7 @@
 {
   "type": "module",
   "name": "@amitdeshmukh/ax-crew",
-  "version": "4.1.2",
+  "version": "6.0.0",
   "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",
@@ -11,7 +11,7 @@
   "scripts": {
     "build": "rm -rf dist && tsc --outDir dist",
     "release": "npm run build && npm publish --access public",
-    "test": "vitest",
+    "test": "vitest run",
     "test:watch": "vitest watch",
     "test:coverage": "vitest run --coverage",
     "test:ui": "vitest --ui"
@@ -24,8 +24,8 @@
     "uuid": "^10.0.0"
   },
   "peerDependencies": {
-    "@ax-llm/ax": "14.0.16",
-    "@ax-llm/ax-tools": "14.0.16"
+    "@ax-llm/ax": "^14.0.36",
+    "@ax-llm/ax-tools": "^14.0.36"
   },
   "devDependencies": {
     "@testing-library/jest-dom": "^6.6.3",

package/src/agents/agentConfig.ts

@@ -1,13 +1,12 @@
-import fs from 'fs';
-// Import Ax factory and MCP transports (as exported by current package)
+// Import Ax factory and MCP transports
 import { ai, AxMCPClient, AxMCPHTTPSSETransport, AxMCPStreambleHTTPTransport, AxDefaultCostTracker } from '@ax-llm/ax'
 import type { AxFunction } from '@ax-llm/ax';
 // STDIO transport from tools package
 import { AxMCPStdioTransport } from '@ax-llm/ax-tools'
-import { PROVIDER_API_KEYS } from '../config/index.js';
+// Resolve env by provided key name
 import type { 
   AgentConfig,
-  CrewConfigInput,
+  AxCrewConfig,
   FunctionRegistryType, 
   MCPTransportConfig, 
   MCPStdioTransportConfig, 
@@ -16,26 +15,6 @@ import type {
 } from '../types.js';
 import type { Provider } from '../types.js';
 
-// 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 lives in src/types.ts
-
 // Type guard to check if config is stdio transport
 export function isStdioTransport(config: MCPTransportConfig): config is MCPStdioTransportConfig {
   return 'command' in config;
@@ -62,53 +41,7 @@ function isConstructor<T>(func: any): func is { new (...args: any[]): T } {
   return typeof func === 'function' && 'prototype' in func && 'toFunction' in func.prototype;
 }
 
-/**
- * Provides a user-friendly error message for JSON parsing errors
- */
-const getFormattedJSONError = (error: Error, fileContents: string): string => {
-  if (error instanceof SyntaxError) {
-    const match = error.message.match(/position (\d+)/);
-    const position = match ? parseInt(match[1]) : -1;
-    
-    if (position !== -1) {
-      const lines = fileContents.split('\n');
-      let currentPos = 0;
-      let errorLine = 0;
-      let errorColumn = 0;
-
-      // Find the line and column of the error
-      for (let i = 0; i < lines.length; i++) {
-        if (currentPos + lines[i].length >= position) {
-          errorLine = i + 1;
-          errorColumn = position - currentPos + 1;
-          break;
-        }
-        currentPos += lines[i].length + 1; // +1 for the newline character
-      }
-
-      const contextLines = lines.slice(Math.max(0, errorLine - 3), errorLine + 2)
-        .map((line, idx) => `${errorLine - 2 + idx}:  ${line}`).join('\n');
-
-      return `JSON Parse Error in your agent configuration:
-      
-Error near line ${errorLine}, column ${errorColumn}
-
-Context:
-${contextLines}
-
-Common issues to check:
-- Missing or extra commas between properties
-- Missing quotes around property names
-- Unmatched brackets or braces
-- Invalid JSON values
-- Trailing commas (not allowed in JSON)
-
-Original error: ${error.message}`;
-    }
-  }
-  
-  return `Error parsing agent configuration: ${error.message}`;
-};
+// Removed file/JSON parse helpers to keep browser-safe
 
 const initializeMCPServers = async (agentConfigData: AgentConfig): Promise<AxFunction[]> => {
   const mcpServers = agentConfigData.mcpServers;
@@ -151,32 +84,16 @@ const initializeMCPServers = async (agentConfigData: AgentConfig): Promise<AxFun
 };
 
 /**
- * Parses and returns the AxCrew config from either a JSON config file or a direct JSON object.
- * @param {CrewConfigInput} input - Either a path to the agent config json file or a JSON object with crew configuration.
+ * Returns the AxCrew config from a direct JSON object. Browser-safe.
+ * @param {CrewConfig} input - A JSON object with crew configuration.
  * @returns {Object} The parsed crew config.
  * @throws Will throw an error if reading/parsing fails.
  */
-const parseCrewConfig = (input: CrewConfigInput): { crew: AgentConfig[] } => {
-  try {
-    if (typeof input === 'string') {
-      // Handle file path input
-      const fileContents = fs.readFileSync(input, 'utf8');
-      const parsedConfig = JSON.parse(fileContents) as { crew: AgentConfig[] };
-      return parsedConfig;
-    } else {
-      // Handle direct JSON object input
-      return input;
-    }
-  } catch (e) {
-    if (e instanceof Error) {
-      if (typeof input === 'string') {
-        const formattedError = getFormattedJSONError(e, fs.readFileSync(input, 'utf8'));
-        throw new Error(formattedError);
-      }
-      throw new Error(`Error parsing agent configuration: ${e.message}`);
-    }
-    throw e;
+const parseCrewConfig = (input: AxCrewConfig): { crew: AgentConfig[] } => {
+  if (!input || typeof input !== 'object' || !Array.isArray((input as any).crew)) {
+    throw new Error('Invalid crew configuration: expected an object with a crew array');
   }
+  return input as { crew: AgentConfig[] };
 };
 
 /**
@@ -184,7 +101,7 @@ const parseCrewConfig = (input: CrewConfigInput): { crew: AgentConfig[] } => {
  * and creates an instance of the Agent with the appropriate settings.
  *
  * @param {string} agentName - The identifier for the AI agent to be initialized.
- * @param {CrewConfigInput} crewConfig - Either a file path to the JSON configuration or a JSON object with crew configuration.
+ * @param {AxCrewConfig} crewConfig - A JSON object with crew configuration.
  * @param {FunctionRegistryType} functions - The functions available to the agent.
  * @param {Object} state - The state object for the agent.
  * @returns {Object} An object containing the Agents AI instance, its name, description, signature, functions and subAgentList.
@@ -193,7 +110,7 @@ const parseCrewConfig = (input: CrewConfigInput): { crew: AgentConfig[] } => {
  */
 const parseAgentConfig = async (
   agentName: string, 
-  crewConfig: CrewConfigInput,
+  crewConfig: AxCrewConfig,
   functions: FunctionRegistryType,
   state: Record<string, any>
 ) => {
@@ -204,20 +121,17 @@ const parseAgentConfig = async (
       throw new Error(`AI agent with name ${agentName} is not configured`);
     }
 
-    // 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(', ')}`);
-    }
+    // Normalize provider slug to lowercase and validate via Ax factory
+    const lower = String(agentConfigData.provider).toLowerCase() as Provider;
     const provider = lower as Provider;
 
-    // If an API Key property is present, get the API key for the AI agent from the environment variables
+    // Resolve API key from user-supplied environment variable name
     let apiKey = '';
     if (agentConfigData.providerKeyName) {
-      apiKey = PROVIDER_API_KEYS[agentConfigData.providerKeyName] || '';
-
+      const keyName = agentConfigData.providerKeyName;
+      apiKey = resolveApiKey(keyName) || '';
       if (!apiKey) {
-        throw new Error(`API key for provider ${agentConfigData.provider} is not set in environment variables`);
+        throw new Error(`API key '${keyName}' for provider ${agentConfigData.provider} is not set in environment`);
       }
     } else {
       throw new Error(`Provider key name is missing in the agent configuration`);
@@ -246,7 +160,19 @@ const parseAgentConfig = async (
         throw new Error(`Invalid apiURL provided: ${agentConfigData.apiURL}`);
       }
     }
-    const aiInstance = ai(aiArgs);
+    // Forward provider-specific arguments as-is; let Ax validate/ignore as needed
+    const providerArgs = (agentConfigData as any).providerArgs;
+    if (providerArgs && typeof providerArgs === 'object') {
+      Object.assign(aiArgs, providerArgs);
+    }
+    // Validate provider by attempting instantiation; Ax will throw on unknown providers
+    let aiInstance;
+    try {
+      aiInstance = ai(aiArgs);
+    } catch (e) {
+      const msg = e instanceof Error ? e.message : String(e);
+      throw new Error(`Unsupported provider '${provider}': ${msg}`);
+    }
 
     // If an mcpServers config is provided in the agent config, convert to functions
     const mcpFunctions = await initializeMCPServers(agentConfigData);
@@ -291,4 +217,21 @@ const parseAgentConfig = async (
 export { 
   parseAgentConfig,
   parseCrewConfig
-};
+};
+
+function resolveApiKey(varName: string): string | undefined {
+  try {
+    // Prefer Node env when available
+    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+    // @ts-ignore
+    if (typeof process !== 'undefined' && process?.env) {
+      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+      // @ts-ignore
+      return process.env[varName];
+    }
+    // Fallback: allow global exposure in browser builds (e.g., injected at runtime)
+    return (globalThis as any)?.[varName];
+  } catch {
+    return undefined;
+  }
+}

package/src/agents/agentUseCosts.ts

@@ -15,6 +15,10 @@ import type {
 export class StateFulAxAgentUsage {
   static STATE_KEY_PREFIX = 'agent_usage_';
 
+  /**
+   * Compute usage costs given a model usage record and model pricing info.
+   * Returns null if inputs are invalid. Token-based costs are computed with high precision.
+   */
   static calculateCost(modelUsage: ModelUsage, modelInfo: ModelInfo): UsageCost | null {
     // Handle both direct properties and nested tokens structure
     const promptTokens = (modelUsage as any).tokens?.promptTokens ?? modelUsage.promptTokens;
@@ -59,6 +63,10 @@ export class StateFulAxAgentUsage {
     }
   }
 
+  /**
+   * Persist or aggregate the cost for an agent in the shared crew state.
+   * No-op if cost is null.
+   */
   static trackCostInState(agentName: string, cost: UsageCost | null, state: StateInstance) {
     // If cost is null, skip tracking
     if (!cost) return;
@@ -90,6 +98,9 @@ export class StateFulAxAgentUsage {
     }
   }
 
+  /**
+   * Aggregate and return total costs across all agents from the shared crew state.
+   */
   static getAggregatedCosts(state: StateInstance): AggregatedCosts {
     const allState = state.getAll();
     const agentCosts: Record<string, UsageCost> = {};
@@ -128,6 +139,9 @@ export class StateFulAxAgentUsage {
     };
   }
 
+  /**
+   * Remove all stored per-agent costs from the shared crew state.
+   */
   static resetCosts(state: StateInstance) {
     const allState = state.getAll();
     Object.keys(allState).forEach(key => {

package/src/agents/compose.ts

@@ -0,0 +1,80 @@
+import { ai } from '@ax-llm/ax';
+import type { AxAI } from '@ax-llm/ax';
+import type { AxCrewConfig } from '../types.js';
+
+type BuildProviderArgs = {
+  provider: string;
+  apiKey: string;
+  config: any;
+  apiURL?: string;
+  providerArgs?: Record<string, unknown>;
+  options?: Record<string, unknown>;
+};
+
+export function instantiateProvider({
+  provider,
+  apiKey,
+  config,
+  apiURL,
+  providerArgs,
+  options,
+}: BuildProviderArgs): AxAI<any> {
+  const args: any = { name: provider as any, apiKey, config, options };
+  if (apiURL) args.apiURL = apiURL;
+  if (providerArgs && typeof providerArgs === 'object') Object.assign(args, providerArgs);
+  return ai(args) as unknown as AxAI<any>;
+}
+
+export function buildProvidersFromConfig(cfg: AxCrewConfig): AxAI<any>[] {
+  const services: AxAI<any>[] = [];
+  for (const agent of cfg.crew) {
+    const apiKeyName = agent.providerKeyName;
+    if (!apiKeyName) throw new Error(`Provider key name is missing for agent ${agent.name}`);
+    const apiKey = resolveApiKey(apiKeyName) || '';
+    if (!apiKey) throw new Error(`API key '${apiKeyName}' not set for agent ${agent.name}`);
+
+    const service = instantiateProvider({
+      provider: String(agent.provider).toLowerCase(),
+      apiKey,
+      config: agent.ai,
+      apiURL: agent.apiURL,
+      providerArgs: (agent as any).providerArgs,
+      options: agent.options,
+    });
+    services.push(service);
+  }
+  return services;
+}
+
+
+// Provider discovery helpers consolidated here (previously in src/providers.ts)
+export function discoverProvidersFromConfig(cfg: AxCrewConfig): string[] {
+  const providers = new Set<string>();
+  for (const agent of cfg.crew) {
+    providers.add(String(agent.provider).toLowerCase());
+  }
+  return Array.from(providers);
+}
+
+export function listSelectableProviders(cfg: AxCrewConfig): string[] {
+  return discoverProvidersFromConfig(cfg);
+}
+
+function resolveApiKey(varName: string): string | undefined {
+  try {
+    // Prefer Node env when available
+    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+    // @ts-ignore
+    if (typeof process !== 'undefined' && process?.env) {
+      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+      // @ts-ignore
+      return process.env[varName];
+    }
+    // Fallback: allow global exposure in browser builds (e.g., injected at runtime)
+    return (globalThis as any)?.[varName];
+  } catch {
+    return undefined;
+  }
+}
+
+

package/src/agents/index.ts

@@ -14,7 +14,7 @@ import type {
    StateInstance, 
    FunctionRegistryType, 
    UsageCost, 
-   CrewConfigInput, 
+   AxCrewConfig,
    MCPTransportConfig,
 } from "../types.js";
 
@@ -238,10 +238,20 @@ class StatefulAxAgent extends AxAgent<any, any> {
   getAccumulatedCosts(): UsageCost | null { return null; }
 
   // Metrics API for this agent
+  /**
+   * 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() {
     const crewId = (this.state as any)?.crewId || (this.state.get?.('crewId')) || 'default';
     return MetricsRegistry.snapshot({ crewId, agent: this.agentName } as any);
   }
+  /**
+   * Reset all tracked metrics for this agent (does not affect other agents).
+   * Call this to start fresh measurement windows for the agent.
+   */
   resetMetrics(): void {
     const crewId = (this.state as any)?.crewId || (this.state.get?.('crewId')) || 'default';
     MetricsRegistry.reset({ crewId, agent: this.agentName } as any);
@@ -250,10 +260,26 @@ class StatefulAxAgent extends AxAgent<any, any> {
 }
 
 /**
- * Represents a crew of agents with shared state functionality.
+ * 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
  */
 class AxCrew {
-  private crewConfig: CrewConfigInput;
+  private crewConfig: AxCrewConfig;
   functionsRegistry: FunctionRegistryType = {};
   crewId: string;
   agents: Map<string, StatefulAxAgent> | null;
@@ -261,12 +287,12 @@ class AxCrew {
 
   /**
    * Creates an instance of AxCrew.
-   * @param {CrewConfigInput} crewConfig - Either a path to the agent config file or a JSON object with crew configuration.
+   * @param {AxCrewConfig} crewConfig - JSON object with crew configuration.
    * @param {FunctionRegistryType} [functionsRegistry={}] - The registry of functions to use in the crew.
    * @param {string} [crewId=uuidv4()] - The unique identifier for the crew.
    */
   constructor(
-    crewConfig: CrewConfigInput,
+    crewConfig: AxCrewConfig,
     functionsRegistry: FunctionRegistryType = {},
     crewId: string = uuidv4()
   ) {
@@ -276,7 +302,7 @@ class AxCrew {
     }
 
     // Validate each agent in the crew
-    crewConfig.crew.forEach(agent => {
+    crewConfig.crew.forEach((agent: any) => {
       if (!agent.name || agent.name.trim() === '') {
         throw new Error('Agent name cannot be empty');
       }
@@ -499,7 +525,8 @@ class AxCrew {
 
 
   /**
-   * Resets all cost tracking for the crew
+   * 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 {
     // Reset AxAgent built-in usage and our metrics registry
@@ -513,9 +540,19 @@ class AxCrew {
   }
 
   // Metrics API
+  /**
+   * 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() {
     return MetricsRegistry.snapshotCrew(this.crewId);
   }
+  /**
+   * Reset all tracked metrics for the entire crew.
+   * Use to clear totals before a new measurement period.
+   */
   resetCrewMetrics(): void {
     MetricsRegistry.reset({ crewId: this.crewId });
   }

package/src/functions/index.ts

@@ -1,6 +1,16 @@
 import { CurrentDateTime, DaysBetweenDates } from './dateTime.js';
 
-// Built-in functions
+/**
+ * Built-in function registry for AxCrew agents.
+ *
+ * Contains common utility tools/functions that can be referenced by name from
+ * agent configs (e.g., "functions": ["CurrentDateTime", "DaysBetweenDates"]).
+ * You can pass this object to the AxCrew constructor or merge with your
+ * own registry.
+ * Example:
+ * const crew = new AxCrew(config, AxCrewFunctions); or
+ * const crew = new AxCrew(config, { ...AxCrewFunctions, ...myFunctions });
+ */
 const AxCrewFunctions = {
   CurrentDateTime,
   DaysBetweenDates

package/src/index.ts

@@ -1,6 +1,6 @@
 import { AxCrew } from './agents/index.js';
 import { AxCrewFunctions } from './functions/index.js';
-import type { CrewConfigInput, AgentConfig } from './types.js';
+import type { AxCrewConfig, AgentConfig } from './types.js';
 
 import type { 
   UsageCost, 
@@ -9,28 +9,42 @@ import type {
   StateInstance, 
   FunctionRegistryType 
 } from './types.js';
+/**
+ * Metrics types and helpers for request counts, token usage, and estimated cost.
+ *
+ * Re-exports the metrics module for convenience:
+ *  - Types: TokenUsage, MetricsSnapshot, etc.
+ *  - Namespace: MetricsRegistry (record/snapshot/reset helpers)
+ */
 export * from './metrics/index.js';
+/**
+ * MetricsRegistry provides functions to record requests, tokens, and cost,
+ * and to snapshot/reset metrics at agent or crew granularity.
+ */
 export { MetricsRegistry } from './metrics/index.js';
 
-// Main AxCrew configuration interface
 /**
- * The configuration for an AxCrew.
- * 
- * @property {AgentConfig[]} crew - The agents that make up the crew.
+ * Create and manage a crew of Ax agents that share state and metrics.
+ * See the `AxCrew` class for full documentation.
+ */
+const _AxCrew: typeof AxCrew = AxCrew;
+
+/**
+ * Built-in function registry with common tools that can be referenced by name
+ * from agent configs, or extended with your own functions.
  */
-interface AxCrewConfig {
-  crew: AgentConfig[];
-}
+const _AxCrewFunctions: typeof AxCrewFunctions = AxCrewFunctions;
 
 export { 
-  AxCrew,
-  AxCrewFunctions,
+  /** 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 exports
   type AggregatedMetrics,
   type AggregatedCosts,
   type AgentConfig,
-  type CrewConfigInput,
   type AxCrewConfig,
   type StateInstance,
   type UsageCost,

package/src/metrics/registry.ts

@@ -44,6 +44,12 @@ function getOrInit(labels: LabelKeys): Counters {
   return c;
 }
 
+/**
+ * Record a completed request.
+ * @param labels Crew/agent/provider/model identifiers
+ * @param streaming Whether this was a streaming request
+ * @param durationMs Duration in milliseconds
+ */
 export function recordRequest(labels: LabelKeys, streaming: boolean, durationMs: number) {
   const c = getOrInit(labels);
   c.requests += 1;
@@ -52,17 +58,26 @@ export function recordRequest(labels: LabelKeys, streaming: boolean, durationMs:
   c.durationCount += 1;
 }
 
+/**
+ * Record an error occurrence for the given labels.
+ */
 export function recordError(labels: LabelKeys) {
   const c = getOrInit(labels);
   c.errors += 1;
 }
 
+/**
+ * Record token usage for a request (prompt and completion).
+ */
 export function recordTokens(labels: LabelKeys, usage: TokenUsage) {
   const c = getOrInit(labels);
   c.inputTokens += usage.promptTokens || 0;
   c.outputTokens += usage.completionTokens || 0;
 }
 
+/**
+ * Add estimated cost (USD) to the cumulative total for the labels.
+ */
 export function recordEstimatedCost(labels: LabelKeys, usd: number) {
   const c = getOrInit(labels);
   const current = new Big(c.estimatedCostUSD || 0);
@@ -70,12 +85,18 @@ export function recordEstimatedCost(labels: LabelKeys, usd: number) {
   c.estimatedCostUSD = Number(current.plus(addition));
 }
 
+/**
+ * Record a function call invocation and add its latency to totals.
+ */
 export function recordFunctionCall(labels: LabelKeys, latencyMs: number) {
   const c = getOrInit(labels);
   c.functionCalls += 1;
   c.functionLatencyMs += latencyMs || 0;
 }
 
+/**
+ * Get a metrics snapshot for specific labels (crew + agent + optional provider/model).
+ */
 export function snapshot(labels: LabelKeys): MetricsSnapshot {
   const c = getOrInit(labels);
   const totalTokens = c.inputTokens + c.outputTokens;
@@ -103,6 +124,9 @@ export function snapshot(labels: LabelKeys): MetricsSnapshot {
   };
 }
 
+/**
+ * Reset metrics for specific labels, or clear all if no labels provided.
+ */
 export function reset(labels?: LabelKeys) {
   if (!labels) {
     store.clear();
@@ -112,6 +136,9 @@ export function reset(labels?: LabelKeys) {
   store.delete(k);
 }
 
+/**
+ * Aggregate a crew-wide metrics snapshot across all agents in the crew.
+ */
 export function snapshotCrew(crewId: string): MetricsSnapshot {
   const empty: Counters = {
     requests: 0,