@amitdeshmukh/ax-crew 5.0.0 → 7.0.0

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/types.ts

@@ -3,26 +3,13 @@ import type {
   AxSignature, 
   AxModelConfig,
   AxMCPStreamableHTTPTransportOptions,
-  AxProgramForwardOptions
+  AxProgramForwardOptions,
+  AxAIArgs
 } from '@ax-llm/ax';
 
-// Canonical provider slugs supported by ai() factory
-export type Provider =
-  // Canonical slugs per docs
-  | 'openai'
-  | 'anthropic'
-  | 'google-gemini'
-  | 'mistral'
-  | 'groq'
-  | 'cohere'
-  | 'together'
-  | 'deepseek'
-  | 'ollama'
-  | 'huggingface'
-  | 'openrouter'
-  | 'azure-openai'
-  | 'reka'
-  | 'x-grok'
+// Provider ids are derived from Ax's factory arg type so new providers added in Ax
+// are picked up at compile time without updating AxCrew.
+export type Provider = AxAIArgs<any>['name'];
 
 /**
  * A state instance that is shared between agents.
@@ -209,13 +196,64 @@ interface AgentConfig {
 }
 
 /**
- * The input type for the agent config. This can be a path to a JSON file or a JSON object.
+ * The configuration object for an AxCrew instance.
+ * 
+ * @property {AgentConfig[]} crew - The agents that make up the crew.
+ * @example
+ * const config: AxCrewConfig = {
+ *   crew: [
+ *     {
+ *       name: "Agent1",
+ *       description: "Agent 1 description",
+ *       signature: "signature",
+ *       provider: "provider",
+ *       providerKeyName: "providerKeyName",
+ *       ai: {
+ *         model: "model",
+ *         temperature: 0,
+ *       },
+ *       options: {
+ *         debug: true,
+ *       },
+ *       functions: ["function1", "function2"],
+ *       agents: ["agent2"],
+ *     },
+ *     {
+ *       name: "Agent2",
+ *       description: "Agent 2 description",
+ *       signature: "signature",
+ *       provider: "provider",
+ *       providerKeyName: "providerKeyName",
+ *       ai: {
+ *         model: "model",
+ *         temperature: 0,
+ *       }
+ *   ]
+ * }
+ * const crew = new AxCrew(config);
  */
-type CrewConfigInput = string | { crew: AgentConfig[] };
+interface AxCrewConfig {
+  crew: AgentConfig[]
+}
+
+/**
+ * Options for the AxCrew instance, specifically allowing optional OpenTelemetry injection.
+ *
+ * @property {Object} [telemetry] - Telemetry configuration.
+ * @property {any} [telemetry.tracer] - OpenTelemetry Tracer instance.
+ * @property {any} [telemetry.meter] - OpenTelemetry Meter instance.
+ */
+interface AxCrewOptions {
+  telemetry?: {
+    tracer?: any;
+    meter?: any;
+  }
+}
 
 export {
   type AgentConfig,
-  type CrewConfigInput,
+  type AxCrewConfig,
+  type AxCrewOptions,
   type AggregatedMetrics,
   type StateInstance,
   type FunctionRegistryType,

package/tests/telemetry.test.ts

@@ -0,0 +1,81 @@
+
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { AxCrew } from '../src/index.js';
+import type { AxCrewConfig } from '../src/types.js';
+import * as axModule from '@ax-llm/ax';
+
+// Mock the entire @ax-llm/ax module
+vi.mock('@ax-llm/ax', async (importOriginal) => {
+  const actual = await importOriginal() as typeof axModule;
+  return {
+    ...actual,
+    // Spy on the ai factory function
+    ai: vi.fn().mockImplementation((args) => {
+      // Return a dummy object that mimics enough of AxAI to satisfy AxCrew
+      return {
+        getName: () => args.name,
+        chat: vi.fn(),
+        defaults: { model: args.config?.model },
+        options: args.options // Store options so we can potentially inspect if needed, though we rely on the spy
+      };
+    }),
+    // We need to keep other exports working if they are used
+    AxAgent: actual.AxAgent,
+    AxAI: actual.AxAI,
+    AxDefaultCostTracker: actual.AxDefaultCostTracker
+  };
+});
+
+describe('AxCrew Telemetry', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it('should pass telemetry options to the underlying AxAI factory', async () => {
+    const mockTracer = { isMockTracer: true };
+    const mockMeter = { isMockMeter: true };
+
+    const crewConfig: AxCrewConfig = {
+      crew: [
+        {
+          name: "telemetry-agent",
+          description: "An agent for testing telemetry propagation",
+          signature: "in:string -> out:string",
+          provider: "openai",
+          providerKeyName: "OPENAI_API_KEY",
+          ai: { model: "gpt-4o-mini" }
+        }
+      ]
+    };
+
+    // Set dummy key
+    process.env.OPENAI_API_KEY = 'dummy-key';
+
+    const options = {
+        telemetry: {
+            tracer: mockTracer,
+            meter: mockMeter
+        }
+    };
+
+    // Initialize Crew with telemetry options (3rd argument now)
+    const crew = new AxCrew(crewConfig, {}, options as any);
+
+    // Create the agent
+    await crew.addAgent("telemetry-agent");
+
+    // Check if the 'ai' mock was called with the expected arguments
+    expect(axModule.ai).toHaveBeenCalled();
+
+    // Get the arguments of the first call to ai()
+    const callArgs = vi.mocked(axModule.ai).mock.calls[0][0];
+
+    // Verify the structure
+    expect(callArgs).toBeDefined();
+    expect(callArgs.options).toBeDefined();
+
+    // Assert tracer and meter were passed correctly
+    expect(callArgs.options.tracer).toBe(mockTracer);
+    expect(callArgs.options.meter).toBe(mockMeter);
+  });
+});

package/dist/config/index.d.ts

@@ -1,5 +0,0 @@
-interface ProviderApiKeys {
-    [key: string]: string | undefined;
-}
-declare const PROVIDER_API_KEYS: ProviderApiKeys;
-export { PROVIDER_API_KEYS, };

package/dist/config/index.js

@@ -1,30 +0,0 @@
-import dotenv from 'dotenv';
-dotenv.config();
-// AI API keys
-const ANTHROPIC_API_KEY = process.env.ANTHROPIC_API_KEY;
-const OPENAI_API_KEY = process.env.OPENAI_API_KEY;
-const AZURE_OPENAI_API_KEY = process.env.AZURE_OPENAI_API_KEY;
-const COHERE_API_KEY = process.env.COHERE_API_KEY;
-const DEEPSEEK_API_KEY = process.env.DEEPSEEK_API_KEY;
-const GEMINI_API_KEY = process.env.GEMINI_API_KEY;
-const GROQ_API_KEY = process.env.GROQ_API_KEY;
-const TOGETHER_API_KEY = process.env.TOGETHER_API_KEY;
-const MISTRAL_API_KEY = process.env.MISTRAL_API_KEY;
-const HUGGINGFACE_API_KEY = process.env.HUGGINGFACE_API_KEY;
-const REKA_API_KEY = process.env.REKA_API_KEY;
-const GROK_API_KEY = process.env.GROK_API_KEY;
-const PROVIDER_API_KEYS = {
-    COHERE_API_KEY,
-    GEMINI_API_KEY,
-    OPENAI_API_KEY,
-    AZURE_OPENAI_API_KEY,
-    ANTHROPIC_API_KEY,
-    DEEPSEEK_API_KEY,
-    GROQ_API_KEY,
-    TOGETHER_API_KEY,
-    MISTRAL_API_KEY,
-    HUGGINGFACE_API_KEY,
-    REKA_API_KEY,
-    GROK_API_KEY
-};
-export { PROVIDER_API_KEYS, };

package/src/config/index.ts

@@ -1,40 +0,0 @@
-import dotenv from 'dotenv';
-dotenv.config();
-
-// AI API keys
-const ANTHROPIC_API_KEY: string | undefined = process.env.ANTHROPIC_API_KEY;
-const OPENAI_API_KEY: string | undefined = process.env.OPENAI_API_KEY;
-const AZURE_OPENAI_API_KEY: string | undefined = process.env.AZURE_OPENAI_API_KEY;
-const COHERE_API_KEY: string | undefined = process.env.COHERE_API_KEY;
-const DEEPSEEK_API_KEY: string | undefined = process.env.DEEPSEEK_API_KEY;
-const GEMINI_API_KEY: string | undefined = process.env.GEMINI_API_KEY;
-const GROQ_API_KEY: string | undefined = process.env.GROQ_API_KEY;
-const TOGETHER_API_KEY: string | undefined = process.env.TOGETHER_API_KEY;
-const MISTRAL_API_KEY: string | undefined = process.env.MISTRAL_API_KEY;
-const HUGGINGFACE_API_KEY: string | undefined = process.env.HUGGINGFACE_API_KEY;
-const REKA_API_KEY: string | undefined = process.env.REKA_API_KEY;
-const GROK_API_KEY: string | undefined = process.env.GROK_API_KEY;
-// Note: Ollama typically doesn't require an API key for local usage
-
-interface ProviderApiKeys {
-  [key: string]: string | undefined;
-}
-
-const PROVIDER_API_KEYS: ProviderApiKeys = {
-  COHERE_API_KEY,
-  GEMINI_API_KEY,
-  OPENAI_API_KEY,
-  AZURE_OPENAI_API_KEY,
-  ANTHROPIC_API_KEY,
-  DEEPSEEK_API_KEY,
-  GROQ_API_KEY,
-  TOGETHER_API_KEY,
-  MISTRAL_API_KEY,
-  HUGGINGFACE_API_KEY,
-  REKA_API_KEY,
-  GROK_API_KEY
-};
-
-export {
-  PROVIDER_API_KEYS,
-};