@amitdeshmukh/ax-crew 3.11.1 → 4.0.1

package/src/metrics/registry.ts

@@ -0,0 +1,167 @@
+import type { LabelKeys, MetricsSnapshot, TokenUsage } from './types.js';
+import Big from 'big.js';
+
+type Key = string;
+
+type Counters = {
+  requests: number;
+  errors: number;
+  streaming: number;
+  durationMsSum: number;
+  durationCount: number;
+  inputTokens: number;
+  outputTokens: number;
+  estimatedCostUSD: number;
+  functionCalls: number;
+  functionLatencyMs: number;
+};
+
+const store = new Map<Key, Counters>();
+
+function keyOf(labels: LabelKeys): Key {
+  const { crewId, agent = '', provider = '', model = '' } = labels;
+  return [crewId, agent, provider, model].join('|');
+}
+
+function getOrInit(labels: LabelKeys): Counters {
+  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: LabelKeys, streaming: boolean, durationMs: number) {
+  const c = getOrInit(labels);
+  c.requests += 1;
+  if (streaming) c.streaming += 1;
+  c.durationMsSum += durationMs;
+  c.durationCount += 1;
+}
+
+export function recordError(labels: LabelKeys) {
+  const c = getOrInit(labels);
+  c.errors += 1;
+}
+
+export function recordTokens(labels: LabelKeys, usage: TokenUsage) {
+  const c = getOrInit(labels);
+  c.inputTokens += usage.promptTokens || 0;
+  c.outputTokens += usage.completionTokens || 0;
+}
+
+export function recordEstimatedCost(labels: LabelKeys, usd: number) {
+  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: LabelKeys, latencyMs: number) {
+  const c = getOrInit(labels);
+  c.functionCalls += 1;
+  c.functionLatencyMs += latencyMs || 0;
+}
+
+export function snapshot(labels: LabelKeys): MetricsSnapshot {
+  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?: LabelKeys) {
+  if (!labels) {
+    store.clear();
+    return;
+  }
+  const k = keyOf(labels);
+  store.delete(k);
+}
+
+export function snapshotCrew(crewId: string): MetricsSnapshot {
+  const empty: Counters = {
+    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,
+    },
+  } as MetricsSnapshot;
+}
+
+

package/src/metrics/types.ts

@@ -0,0 +1,48 @@
+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>; // USD per 1K tokens override
+}
+
+

package/src/types.ts

@@ -2,10 +2,27 @@ import type {
   AxFunction, 
   AxSignature, 
   AxModelConfig,
-  AxMCPStreamableHTTPTransportOptions
+  AxMCPStreamableHTTPTransportOptions,
+  AxProgramForwardOptions
 } from '@ax-llm/ax';
 
-import type { Provider } from './agents/agentConfig.js';
+// 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'
 
 /**
  * A state instance that is shared between agents.
@@ -131,7 +148,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
 }
@@ -141,7 +158,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.
@@ -154,7 +171,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 {Record<string, any>} options - Agent options. 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.
@@ -169,7 +186,7 @@ interface AgentConfig {
   ai: AxModelConfig & { model: string };
   debug?: boolean;
   apiURL?: string;
-  options?: Record<string, any>;
+  options?: Partial<AxProgramForwardOptions<any>> & Record<string, any>;
   functions?: string[];
   agents?: string[];
   examples?: Array<Record<string, any>>;
@@ -189,10 +206,10 @@ export {
   type FunctionRegistryType,
   type MCPStdioTransportConfig,
   type MCPHTTPSSETransportConfig,
-  type MCPStreambleHTTPTransportConfig,
+  type MCPStreamableHTTPTransportConfig,
   type MCPTransportConfig,
   type ModelUsage,
   type ModelInfo,
   type UsageCost,
   type AggregatedCosts
-};
+}

package/tests/AxCrew.test.ts

@@ -1,6 +1,10 @@
 import { describe, test, expect, beforeEach } from 'vitest';
 import { AxCrew } from '../src/agents';
 import { AxCrewFunctions } from '../src/functions';
+import type { AxCrewConfig } from '../src/index.js';
+
+// Provide dummy API key for tests
+process.env.ANTHROPIC_API_KEY = process.env.ANTHROPIC_API_KEY || 'test';
 
 describe('AxCrew', () => {
   // Basic initialization tests
@@ -10,7 +14,7 @@ describe('AxCrew', () => {
         crew: [{
           name: "ResearchAgent",
           description: "A research agent that can search the web for information and provide detailed analysis of search results with proper citations",
-          signature: "input:string -> output:string",
+          signature: "query:string -> queryResponse:string",
           provider: "anthropic",
           providerKeyName: "ANTHROPIC_API_KEY",
           ai: {
@@ -19,7 +23,7 @@ describe('AxCrew', () => {
         }]
       };
       
-      const crew = new AxCrew(config, AxCrewFunctions);
+      const crew = new AxCrew(config as AxCrewConfig, AxCrewFunctions);
       expect(crew).toBeInstanceOf(AxCrew);
     });
 
@@ -36,7 +40,7 @@ describe('AxCrew', () => {
           }
         }]
       };
-      expect(() => new AxCrew(invalidConfig, AxCrewFunctions))
+      expect(() => new AxCrew(invalidConfig as AxCrewConfig, AxCrewFunctions))
         .toThrowError('Agent name cannot be empty');
     });
   });
@@ -50,8 +54,8 @@ describe('AxCrew', () => {
         crew: [
           {
             name: "agent1",
-            description: "A sophisticated agent that processes text input and generates structured analysis with detailed explanations and recommendations",
-            signature: "input:string -> output:string",
+            description: "A sophisticated agent that processes text query and generates structured analysis with detailed explanations and recommendations",
+            signature: "query:string -> queryResponse:string",
             provider: "anthropic",
             providerKeyName: "ANTHROPIC_API_KEY",
             ai: { model: "claude-3-haiku-20240307" }
@@ -59,7 +63,7 @@ describe('AxCrew', () => {
           {
             name: "agent2",
             description: "An advanced processing agent that builds upon agent1's output to provide deeper insights and actionable intelligence",
-            signature: "input:string -> output:string",
+            signature: "query:string -> queryResponse:string",
             provider: "anthropic",
             providerKeyName: "ANTHROPIC_API_KEY",
             ai: { model: "claude-3-haiku-20240307" },
@@ -94,7 +98,7 @@ describe('AxCrew', () => {
         crew: [{
           name: "testAgent",
           description: "A comprehensive testing agent that validates inputs, processes data, and ensures output quality through multiple verification steps",
-          signature: "input:string -> output:string",
+          signature: "query:string -> queryResponse:string",
           provider: "anthropic",
           providerKeyName: "ANTHROPIC_API_KEY",
           ai: { model: "claude-3-haiku-20240307" }
@@ -104,19 +108,19 @@ describe('AxCrew', () => {
       await crew.addAgent('testAgent');
     });
 
-    test('should track costs correctly', async () => {
-      const costs = crew.getAggregatedCosts();
-      expect(costs).toBeDefined();
-      expect(costs).toHaveProperty('totalCost');
-      expect(typeof costs.totalCost).toBe('string');
+    test('should track costs correctly (metrics)', async () => {
+      const metrics = crew.getCrewMetrics();
+      expect(metrics).toBeDefined();
+      expect(metrics).toHaveProperty('estimatedCostUSD');
+      expect(typeof metrics.estimatedCostUSD).toBe('number');
     });
 
-    test('should reset costs', () => {
+    test('should reset costs (metrics)', () => {
       crew.resetCosts();
-      const costs = crew.getAggregatedCosts();
-      expect(costs).toBeDefined();
-      expect(costs).toHaveProperty('totalCost');
-      expect(costs.totalCost).toBe('0');  // Updated to match actual format
+      const metrics = crew.getCrewMetrics();
+      expect(metrics).toBeDefined();
+      expect(metrics).toHaveProperty('estimatedCostUSD');
+      expect(metrics.estimatedCostUSD).toBe(0);
     });
   });
 
@@ -127,7 +131,7 @@ describe('AxCrew', () => {
         crew: [{
           name: "testAgent",
           description: "A comprehensive testing agent that validates inputs, processes data, and ensures output quality through multiple verification steps",
-          signature: "input:string -> output:string",
+          signature: "query:string -> queryResponse:string",
           provider: "anthropic",
           providerKeyName: "ANTHROPIC_API_KEY",
           ai: { model: "claude-3-haiku-20240307" }