@amitdeshmukh/ax-crew 8.0.3 → 8.1.0

package/tests/execution-mode-metrics.test.ts

@@ -0,0 +1,231 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import type { AxCrewConfig } from "../src/types.js";
+import { AxCrew, MetricsRegistry } from "../src/index.js";
+import * as axModule from "@ax-llm/ax";
+
+vi.mock("@ax-llm/ax", async (importOriginal) => {
+  const actual = (await importOriginal()) as typeof axModule;
+
+  class MockAxDefaultCostTracker {
+    private total = 0;
+
+    trackTokens(count: number) {
+      this.total += (count || 0) * 0.000001;
+    }
+
+    getCurrentCost() {
+      return this.total;
+    }
+  }
+
+  class MockAxGen {
+    private usage: any[] = [];
+    private signature: any;
+
+    constructor(signature: any) {
+      this.signature = signature;
+    }
+
+    async forward(_ai: any, values: Record<string, any>) {
+      this.usage.push({
+        ai: "mock-ai",
+        model: "mock-model",
+        tokens: { promptTokens: 5, completionTokens: 3 },
+      });
+      return { answer: `gen:${values.query ?? values.task ?? "ok"}` };
+    }
+
+    async *streamingForward(_ai: any, values: Record<string, any>) {
+      this.usage.push({
+        ai: "mock-ai",
+        model: "mock-model",
+        tokens: { promptTokens: 5, completionTokens: 3 },
+      });
+      yield { version: 1, index: 0, delta: { answer: `gen:${values.query ?? "ok"}` } };
+    }
+
+    register() {}
+
+    setExamples() {}
+
+    setDescription() {}
+
+    getUsage() {
+      return this.usage;
+    }
+
+    resetUsage() {
+      this.usage = [];
+    }
+
+    getSignature() {
+      return {
+        getDescription: () => "",
+      };
+    }
+  }
+
+  class MockAxAgent {
+    private usage: any[] = [];
+    private signature: any;
+    public program: any;
+    public actorDescription = "";
+    public responderDescription = "";
+
+    constructor(config: any) {
+      this.signature = {
+        getDescription: () => "",
+      };
+      this.program = {
+        setDescription: vi.fn(),
+      };
+      this.actorDescription = config?.agentIdentity?.description ?? "";
+      this.responderDescription = config?.agentIdentity?.description ?? "";
+    }
+
+    async forward(_ai: any, values: Record<string, any>) {
+      this.usage.push({
+        ai: "mock-ai",
+        model: "mock-model",
+        tokens: { promptTokens: 11, completionTokens: 7 },
+      });
+      return { answer: `agent:${values.query ?? values.task ?? "ok"}` };
+    }
+
+    async *streamingForward(_ai: any, values: Record<string, any>) {
+      this.usage.push({
+        ai: "mock-ai",
+        model: "mock-model",
+        tokens: { promptTokens: 11, completionTokens: 7 },
+      });
+      yield { version: 1, index: 0, delta: { answer: `agent:${values.query ?? "ok"}` } };
+    }
+
+    getUsage() {
+      return this.usage;
+    }
+
+    resetUsage() {
+      this.usage = [];
+    }
+
+    getFunction() {
+      return { name: "mockAgentFn", description: "mock", parameters: { type: "object", properties: {} } };
+    }
+
+    getSignature() {
+      return this.signature;
+    }
+
+    _buildSplitPrograms() {}
+  }
+
+  return {
+    ...actual,
+    ai: vi.fn().mockImplementation((args: any) => ({
+      getName: () => args.name,
+      chat: vi.fn(),
+      defaults: { model: args.config?.model ?? "mock-model" },
+      options: args.options,
+    })),
+    AxAgent: MockAxAgent as any,
+    AxGen: MockAxGen as any,
+    AxDefaultCostTracker: MockAxDefaultCostTracker as any,
+  };
+});
+
+const makeConfig = (mode: "axagent" | "axgen"): AxCrewConfig => ({
+  crew: [
+    {
+      name: `mode-${mode}`,
+      description: "Agent for execution mode metrics testing",
+      executionMode: mode,
+      signature: "query:string -> answer:string",
+      provider: "openai" as any,
+      providerKeyName: "OPENAI_API_KEY",
+      ai: { model: "mock-model" } as any,
+      axAgentOptions:
+        mode === "axagent"
+          ? {
+              contextFields: [],
+              mode: "simple",
+            }
+          : undefined,
+    },
+  ],
+});
+
+describe("Execution Mode Metrics + Tracing", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    MetricsRegistry.reset();
+    process.env.OPENAI_API_KEY = "dummy-key";
+  });
+
+  afterEach(() => {
+    MetricsRegistry.reset();
+  });
+
+  it("records built-in usage/cost and telemetry wiring in axagent mode", async () => {
+    const telemetry = { tracer: { id: "tracer" }, meter: { id: "meter" } };
+    const crew = new AxCrew(makeConfig("axagent"), {}, { telemetry } as any);
+    await crew.addAllAgents();
+
+    const agent = crew.agents?.get("mode-axagent");
+    expect(agent).toBeDefined();
+
+    const result = await agent!.forward({ query: "hello" });
+    expect(result.answer).toContain("agent:");
+
+    const aiCallArgs = vi.mocked(axModule.ai).mock.calls[0][0] as any;
+    expect(aiCallArgs.options.tracer).toBe(telemetry.tracer);
+    expect(aiCallArgs.options.meter).toBe(telemetry.meter);
+
+    const crewMetrics = crew.getCrewMetrics();
+    expect(crewMetrics.requests.totalRequests).toBe(1);
+    expect(crewMetrics.tokens.totalTokens).toBe(18); // 11 + 7 from MockAxAgent
+    expect(crewMetrics.estimatedCostUSD).toBeGreaterThan(0);
+  });
+
+  it("records built-in usage/cost and telemetry wiring in axgen mode", async () => {
+    const telemetry = { tracer: { id: "tracer" }, meter: { id: "meter" } };
+    const crew = new AxCrew(makeConfig("axgen"), {}, { telemetry } as any);
+    await crew.addAllAgents();
+
+    const agent = crew.agents?.get("mode-axgen");
+    expect(agent).toBeDefined();
+
+    const result = await agent!.forward({ query: "hello" });
+    expect(result.answer).toContain("gen:");
+
+    const aiCallArgs = vi.mocked(axModule.ai).mock.calls[0][0] as any;
+    expect(aiCallArgs.options.tracer).toBe(telemetry.tracer);
+    expect(aiCallArgs.options.meter).toBe(telemetry.meter);
+
+    const crewMetrics = crew.getCrewMetrics();
+    expect(crewMetrics.requests.totalRequests).toBe(1);
+    expect(crewMetrics.tokens.totalTokens).toBe(8); // 5 + 3 from MockAxGen
+    expect(crewMetrics.estimatedCostUSD).toBeGreaterThan(0);
+  });
+
+  it("tracks streaming requests in both execution modes", async () => {
+    for (const mode of ["axagent", "axgen"] as const) {
+      MetricsRegistry.reset();
+      const crew = new AxCrew(makeConfig(mode), {}, { telemetry: {} } as any);
+      await crew.addAllAgents();
+      const agent = crew.agents?.get(`mode-${mode}`);
+      expect(agent).toBeDefined();
+
+      const stream = agent!.streamingForward({ query: "streaming" });
+      for await (const _chunk of stream) {
+        // consume
+      }
+
+      const crewMetrics = crew.getCrewMetrics();
+      expect(crewMetrics.requests.totalRequests).toBe(1);
+      expect(crewMetrics.requests.totalStreamingRequests).toBe(1);
+      expect(crewMetrics.tokens.totalTokens).toBe(mode === "axagent" ? 18 : 8);
+      expect(crewMetrics.estimatedCostUSD).toBeGreaterThan(0);
+    }
+  });
+});

package/examples/factory.ts

@@ -1,174 +0,0 @@
-// Crew Factory - Manages pool of crew instances for concurrent requests
-
-import { AxCrew } from '../dist/index.js';
-import type { AxCrewConfig } from '../dist/types.js';
-import type { AxFunction } from '@ax-llm/ax';
-
-interface CrewInstance {
-  crew: InstanceType<typeof AxCrew>;
-  busy: boolean;
-  id: number;
-  createdAt: number;
-}
-
-interface FactoryConfig {
-  crewConfig: AxCrewConfig;
-  crewFunctions: Record<string, AxFunction>;
-  agentNames: string[];
-}
-
-const MAX_POOL_SIZE = 5;
-const INSTANCE_TTL_MS = 5 * 60 * 1000; // 5 minutes idle before cleanup
-
-let instanceIdCounter = 0;
-const pool: CrewInstance[] = [];
-let isInitialized = false;
-let factoryConfig: FactoryConfig;
-
-/**
- * Initialize the factory with at least one crew instance
- */
-export async function initCrewFactory(config: FactoryConfig): Promise<void> {
-  if (isInitialized) {
-    console.log('[CrewFactory] Already initialized');
-    return;
-  }
-
-  factoryConfig = config;
-
-  console.log('[CrewFactory] Initializing...');
-
-  // Create initial instance
-  await createCrewInstance();
-  isInitialized = true;
-
-  // Start cleanup interval
-  setInterval(cleanupIdleInstances, 60000);
-
-  console.log('[CrewFactory] Initialized with 1 instance');
-}
-
-/**
- * Create a new crew instance
- */
-async function createCrewInstance(): Promise<CrewInstance> {
-  const id = ++instanceIdCounter;
-  console.log(`[CrewFactory] Creating instance #${id}...`);
-
-  const crew = new AxCrew(factoryConfig.crewConfig, factoryConfig.crewFunctions);
-  for (const name of factoryConfig.agentNames) {
-    await crew.addAgentsToCrew([name]);
-  }
-
-  const instance: CrewInstance = {
-    crew,
-    busy: false,
-    id,
-    createdAt: Date.now(),
-  };
-
-  pool.push(instance);
-  console.log(`[CrewFactory] Instance #${id} ready. Pool size: ${pool.length}`);
-
-  return instance;
-}
-
-/**
- * Acquire an available crew instance, creating one if needed
- */
-export async function acquireCrew(): Promise<{ crew: InstanceType<typeof AxCrew>; release: () => void }> {
-  if (!isInitialized) {
-    throw new Error('CrewFactory not initialized. Call initCrewFactory() first.');
-  }
-
-  // Find an available instance
-  let instance = pool.find(i => !i.busy);
-
-  // If none available and pool not at max, create new instance
-  if (!instance && pool.length < MAX_POOL_SIZE) {
-    instance = await createCrewInstance();
-  }
-
-  // If still none available, wait for one to become free
-  if (!instance) {
-    console.log('[CrewFactory] Pool exhausted, waiting for available instance...');
-    instance = await waitForAvailableInstance();
-  }
-
-  instance.busy = true;
-  const instanceId = instance.id;
-
-  console.log(`[CrewFactory] Acquired instance #${instanceId}. Active: ${pool.filter(i => i.busy).length}/${pool.length}`);
-
-  return {
-    crew: instance.crew,
-    release: () => {
-      const inst = pool.find(i => i.id === instanceId);
-      if (inst) {
-        inst.busy = false;
-        inst.createdAt = Date.now(); // Reset TTL on release
-        console.log(`[CrewFactory] Released instance #${instanceId}. Active: ${pool.filter(i => i.busy).length}/${pool.length}`);
-      }
-    },
-  };
-}
-
-/**
- * Wait for an instance to become available
- */
-function waitForAvailableInstance(): Promise<CrewInstance> {
-  return new Promise((resolve) => {
-    const check = () => {
-      const instance = pool.find(i => !i.busy);
-      if (instance) {
-        resolve(instance);
-      } else {
-        setTimeout(check, 100);
-      }
-    };
-    check();
-  });
-}
-
-/**
- * Clean up idle instances beyond the first one
- */
-function cleanupIdleInstances(): void {
-  const now = Date.now();
-  const toRemove: number[] = [];
-
-  for (const instance of pool) {
-    // Keep at least one instance, only cleanup idle instances past TTL
-    if (pool.length - toRemove.length > 1 &&
-        !instance.busy &&
-        now - instance.createdAt > INSTANCE_TTL_MS) {
-      toRemove.push(instance.id);
-    }
-  }
-
-  for (const id of toRemove) {
-    const index = pool.findIndex(i => i.id === id);
-    if (index !== -1) {
-      pool.splice(index, 1);
-      console.log(`[CrewFactory] Cleaned up idle instance #${id}. Pool size: ${pool.length}`);
-    }
-  }
-}
-
-/**
- * Check if factory is ready
- */
-export function isFactoryReady(): boolean {
-  return isInitialized;
-}
-
-/**
- * Get factory stats
- */
-export function getFactoryStats() {
-  return {
-    poolSize: pool.length,
-    activeCount: pool.filter(i => i.busy).length,
-    maxPoolSize: MAX_POOL_SIZE,
-  };
-}

package/examples/google-gemini-test.ts

@@ -1,291 +0,0 @@
-// Standalone test: json[] signature pattern with Google Gemini provider
-//
-// Tests that RoutePlanner returns waypoints as actual objects (json[]),
-// not stringified JSON, after the Gemini schema type-union fix.
-
-import type { AxCrewConfig } from '../dist/types.js';
-import type { AxFunction } from '@ax-llm/ax';
-import { FlightConstantSpeedPlanner, type LLA } from './planner.js';
-import { initCrewFactory, acquireCrew } from './factory.js';
-import dotenv from 'dotenv';
-dotenv.config();
-
-// ---------------------------------------------------------------------------
-// Crew Configuration
-// ---------------------------------------------------------------------------
-
-const crewConfig: AxCrewConfig = {
-  crew: [
-    {
-      name: 'RoutePlanner',
-      description: `
-        Plans optimal flight routes between geolocations.
-        Uses the PlanFlightRoute function to generate smooth paths with fly-by turns.
-        Call PlanFlightRoute with current position and destination to get detailed waypoints.
-      `,
-      signature: `
-        currentLat:number "current latitude in degrees",
-        currentLon:number "current longitude in degrees",
-        currentAlt:number "current altitude in meters",
-        heading?:number "current heading in degrees",
-        speedKts?:number "speed in knots",
-        destLat:number "destination latitude in degrees",
-        destLon:number "destination longitude in degrees",
-        destAlt:number "destination altitude in meters",
-        vehicleType?:string "type of vehicle"
-        ->
-        waypoints:json[] "Array of {lat, lon, alt, heading}",
-        totalDistanceKm:number "total distance in kilometers",
-        eta:string "estimated time of arrival in MM:SS format"
-      `,
-      provider: 'openrouter' as const,
-      providerKeyName: 'OPENROUTER_API_KEY',
-      ai: {
-        model: 'stepfun/step-3.5-flash:free',
-        temperature: 0,
-      },
-      options: {
-        debug: true,
-      },
-      functions: ['PlanFlightRoute'],
-    },
-    {
-      name: 'Manager',
-      description: `Command center assistant that handles user requests for route planning.
-        IMPORTANT: When calling RoutePlanner, you MUST include the waypoints array it returns in your response output.`,
-      signature: `
-        userCommand:string "user request or command"
-        ->
-        responseToCommand:string "A response to the command given.",
-        waypoints?:json[] "Array of {lat, lon, alt, heading} from RoutePlanner",
-        totalDistanceKm?:number "total distance in km from RoutePlanner",
-        eta?:string "ETA from RoutePlanner"
-      `,
-      provider: 'openrouter' as const,
-      providerKeyName: 'OPENROUTER_API_KEY',
-      ai: {
-        model: 'stepfun/step-3.5-flash:free',
-        temperature: 0,
-      },
-      options: {
-        debug: true,
-      },
-      agents: ['RoutePlanner'],
-    },
-  ],
-};
-
-// ---------------------------------------------------------------------------
-// Haversine distance (for dynamic Hz calculation)
-// ---------------------------------------------------------------------------
-
-const EARTH_RADIUS = 6378137.0;
-const MAX_WAYPOINTS = 25;
-const DEFAULT_HZ = 0.0333;
-
-function getDistance(lat1: number, lon1: number, lat2: number, lon2: number): number {
-  const dLat = (lat2 - lat1) * (Math.PI / 180);
-  const dLon = (lon2 - lon1) * (Math.PI / 180);
-  const a =
-    Math.sin(dLat / 2) ** 2 +
-    Math.cos(lat1 * (Math.PI / 180)) *
-      Math.cos(lat2 * (Math.PI / 180)) *
-      Math.sin(dLon / 2) ** 2;
-  return 2 * Math.atan2(Math.sqrt(a), Math.sqrt(1 - a)) * EARTH_RADIUS;
-}
-
-// ---------------------------------------------------------------------------
-// PlanFlightRoute function (uses FlightConstantSpeedPlanner)
-// ---------------------------------------------------------------------------
-
-const planFlightRoute: AxFunction = {
-  name: 'PlanFlightRoute',
-  description: `Plans a flight route from current position through waypoints to destination.
-Uses constant speed with smooth fly-by turns at 25-degree bank angle.
-Returns detailed waypoints along the path plus mission summary (distance, time, ETA).`,
-  parameters: {
-    type: 'object',
-    properties: {
-      currentLat: { type: 'number', description: 'Current latitude in degrees' },
-      currentLon: { type: 'number', description: 'Current longitude in degrees' },
-      currentAlt: { type: 'number', description: 'Current altitude in meters' },
-      heading: { type: 'number', description: 'Current heading in degrees (0-360)' },
-      destLat: { type: 'number', description: 'Destination latitude in degrees' },
-      destLon: { type: 'number', description: 'Destination longitude in degrees' },
-      destAlt: { type: 'number', description: 'Destination altitude in meters' },
-      speedKts: { type: 'number', description: 'Speed in knots' },
-    },
-    required: [
-      'currentLat', 'currentLon', 'currentAlt', 'heading',
-      'destLat', 'destLon', 'destAlt', 'speedKts',
-    ],
-  },
-  func: async (args) => {
-    const {
-      currentLat, currentLon, currentAlt, heading,
-      destLat, destLon, destAlt, speedKts,
-    } = args as Record<string, number>;
-
-    // Convert knots to meters per second (1 knot = 0.514444 m/s)
-    const speedMps = speedKts * 0.514444;
-
-    // Calculate dynamic Hz to limit waypoints for long routes
-    const estimatedDistanceM = getDistance(currentLat, currentLon, destLat, destLon);
-    const estimatedTimeS = estimatedDistanceM / speedMps;
-    const maxHz = MAX_WAYPOINTS / estimatedTimeS;
-    const hz = Math.min(DEFAULT_HZ, maxHz);
-
-    // Create waypoints for the planner
-    const waypoints: LLA[] = [
-      { lat: currentLat, lon: currentLon, alt: currentAlt, heading },
-      { lat: destLat, lon: destLon, alt: destAlt, heading: 0 },
-    ];
-
-    const planner = new FlightConstantSpeedPlanner(speedMps, hz);
-    const { path, summary } = planner.getMissionDetails(waypoints);
-
-    // Return actual objects — tests that json[] works with Gemini
-    const routeWaypoints = path.map((p) => ({
-      lat: Number(p.lat.toFixed(6)),
-      lon: Number(p.lon.toFixed(6)),
-      alt: Math.round(p.alt),
-      heading: p.heading,
-    }));
-
-    return {
-      waypoints: routeWaypoints,
-      totalDistanceKm: summary.totalDistanceKm,
-      totalTimeSeconds: summary.totalTimeSeconds,
-      etaFormatted: summary.etaFormatted,
-      waypointCount: routeWaypoints.length,
-    };
-  },
-};
-
-// ---------------------------------------------------------------------------
-// Function registry
-// ---------------------------------------------------------------------------
-
-const crewFunctions: Record<string, AxFunction> = {
-  PlanFlightRoute: planFlightRoute,
-};
-
-// ---------------------------------------------------------------------------
-// Waypoint validation helpers
-// ---------------------------------------------------------------------------
-
-interface Waypoint {
-  lon: number;
-  lat: number;
-  alt: number;
-}
-
-function normalizeWaypoint(wp: unknown): Waypoint | null {
-  let obj: Record<string, unknown> | null = null;
-
-  // Handle string waypoints (fallback if Gemini returns strings)
-  if (typeof wp === 'string') {
-    try {
-      obj = JSON.parse(wp);
-    } catch {
-      console.log('[Test] Failed to parse waypoint string:', wp);
-      return null;
-    }
-  } else if (typeof wp === 'object' && wp !== null) {
-    obj = wp as Record<string, unknown>;
-  }
-
-  if (!obj) return null;
-
-  const lon =
-    typeof obj.lon === 'number' ? obj.lon :
-    typeof obj.lng === 'number' ? obj.lng : null;
-  const lat = typeof obj.lat === 'number' ? obj.lat : null;
-  const alt =
-    typeof obj.alt === 'number' ? obj.alt :
-    typeof obj.altitude === 'number' ? obj.altitude : 1000;
-
-  if (lon === null || lat === null) return null;
-  if (lat === 0 && lon === 0) return null;
-  if (lat < -90 || lat > 90 || lon < -180 || lon > 180) return null;
-
-  return { lon, lat, alt };
-}
-
-function filterValidWaypoints(waypoints: unknown[]): Waypoint[] {
-  return waypoints
-    .map(normalizeWaypoint)
-    .filter((wp): wp is Waypoint => wp !== null);
-}
-
-// ---------------------------------------------------------------------------
-// Main
-// ---------------------------------------------------------------------------
-
-async function main() {
-  console.log('=== Google Gemini json[] Test ===\n');
-
-  if (!process.env.GEMINI_API_KEY) {
-    console.error('GEMINI_API_KEY not set. Add it to .env or export it.');
-    process.exit(1);
-  }
-
-  // Initialize the factory with our config
-  await initCrewFactory({
-    crewConfig,
-    crewFunctions,
-    agentNames: ['RoutePlanner', 'Manager'],
-  });
-
-  const { crew, release } = await acquireCrew();
-
-  const manager = crew.agents?.get('Manager');
-  if (!manager) {
-    throw new Error('Manager agent not found');
-  }
-
-  try {
-    console.log('[Test] Sending route planning request...\n');
-
-    const result = await manager.forward({
-      userCommand:
-        'Plan a flight route from San Francisco (37.7749, -122.4194, alt 500m) ' +
-        'to Los Angeles (34.0522, -118.2437, alt 500m) at 250 knots, heading 150 degrees.',
-    });
-
-    console.log('\n[Test] Raw Manager result:', JSON.stringify(result, null, 2));
-
-    // Validate response
-    console.log('\n--- Validation ---');
-    console.log('responseToCommand:', result.responseToCommand);
-    console.log('totalDistanceKm:', result.totalDistanceKm);
-    console.log('eta:', result.eta);
-
-    if (Array.isArray(result.waypoints)) {
-      console.log(`\nWaypoints received: ${result.waypoints.length}`);
-
-      // Check if waypoints are objects or strings
-      const first = result.waypoints[0];
-      console.log('First waypoint type:', typeof first);
-      console.log('First waypoint:', JSON.stringify(first));
-
-      const valid = filterValidWaypoints(result.waypoints);
-      console.log(`Valid waypoints: ${valid.length}/${result.waypoints.length}`);
-
-      if (valid.length > 0) {
-        console.log('\n[PASS] json[] pattern worked with Google Gemini');
-      } else {
-        console.log('\n[FAIL] Waypoints received but none were valid');
-      }
-    } else {
-      console.log('\n[FAIL] No waypoints array in response');
-    }
-  } finally {
-    release();
-  }
-}
-
-main().catch((err) => {
-  console.error('[FATAL]', err);
-  process.exit(1);
-});