@amitdeshmukh/ax-crew 3.7.1 → 3.8.1

package/CHANGELOG.md

@@ -5,6 +5,16 @@ This Changelog format is based on [Keep a Changelog]
 adheres to [Semantic Versioning](https://semver.org/spec/
 v2.0.0.html).
 
+## [3.8.1] - 2024-03-28
+
+### Added
+- Enhanced type definitions for model usage metrics with nested token structure support
+
+### Changed
+- Improved error handling and robustness in cost calculations
+- Better null handling in usage cost tracking
+- Updated type definitions in `src/types.ts` to support both direct and nested token structures
+
 ## [3.7.1] - 2025-03-27
 
 ### Fixed

package/dist/agents/agentUseCosts.d.ts

@@ -5,8 +5,8 @@ import type { StateInstance, ModelUsage, ModelInfo, UsageCost, AggregatedCosts }
  */
 export declare class StateFulAxAgentUsage {
     static STATE_KEY_PREFIX: string;
-    static calculateCost(modelUsage: ModelUsage, modelInfo: ModelInfo): UsageCost;
-    static trackCostInState(agentName: string, cost: UsageCost, state: StateInstance): void;
+    static calculateCost(modelUsage: ModelUsage, modelInfo: ModelInfo): UsageCost | null;
+    static trackCostInState(agentName: string, cost: UsageCost | null, state: StateInstance): void;
     static getAggregatedCosts(state: StateInstance): AggregatedCosts;
     static resetCosts(state: StateInstance): void;
 }

package/dist/agents/agentUseCosts.js

@@ -5,45 +5,69 @@ import { Decimal } from 'decimal.js';
  */
 export class StateFulAxAgentUsage {
     static calculateCost(modelUsage, modelInfo) {
-        const { promptTokens, completionTokens } = modelUsage;
+        // Handle both direct properties and nested tokens structure
+        const promptTokens = modelUsage.tokens?.promptTokens ?? modelUsage.promptTokens;
+        const completionTokens = modelUsage.tokens?.completionTokens ?? modelUsage.completionTokens;
         const { promptTokenCostPer1M, completionTokenCostPer1M } = modelInfo;
-        // Use Decimal for precise calculations
-        const promptCost = new Decimal(promptTokens)
-            .div(1000000)
-            .mul(promptTokenCostPer1M)
-            .toDP(10); // Keep 10 decimal places
-        const completionCost = new Decimal(completionTokens)
-            .div(1000000)
-            .mul(completionTokenCostPer1M)
-            .toDP(10);
-        const totalCost = promptCost.plus(completionCost).toDP(10);
-        return {
-            promptCost: promptCost.toString(),
-            completionCost: completionCost.toString(),
-            totalCost: totalCost.toString(),
-            tokenMetrics: {
-                promptTokens,
-                completionTokens,
-                totalTokens: promptTokens + completionTokens
-            }
-        };
+        // Return null instead of throwing errors for invalid values
+        if (typeof promptTokens !== 'number' || isNaN(promptTokens) ||
+            typeof completionTokens !== 'number' || isNaN(completionTokens) ||
+            typeof promptTokenCostPer1M !== 'number' || isNaN(promptTokenCostPer1M) ||
+            typeof completionTokenCostPer1M !== 'number' || isNaN(completionTokenCostPer1M)) {
+            return null;
+        }
+        try {
+            // Use Decimal for precise calculations
+            const promptCost = new Decimal(promptTokens)
+                .div(1000000)
+                .mul(promptTokenCostPer1M)
+                .toDP(10); // Keep 10 decimal places
+            const completionCost = new Decimal(completionTokens)
+                .div(1000000)
+                .mul(completionTokenCostPer1M)
+                .toDP(10);
+            const totalCost = promptCost.plus(completionCost).toDP(10);
+            return {
+                promptCost: promptCost.toString(),
+                completionCost: completionCost.toString(),
+                totalCost: totalCost.toString(),
+                tokenMetrics: {
+                    promptTokens,
+                    completionTokens,
+                    totalTokens: promptTokens + completionTokens
+                }
+            };
+        }
+        catch (error) {
+            // If any decimal calculation fails, return null instead of throwing
+            return null;
+        }
     }
     static trackCostInState(agentName, cost, state) {
+        // If cost is null, skip tracking
+        if (!cost)
+            return;
         const stateKey = `${this.STATE_KEY_PREFIX}${agentName}`;
         const existingCost = state.get(stateKey);
         if (existingCost) {
-            // Aggregate with existing cost using Decimal
-            const aggregatedCost = {
-                promptCost: new Decimal(existingCost.promptCost).plus(cost.promptCost).toDP(10).toString(),
-                completionCost: new Decimal(existingCost.completionCost).plus(cost.completionCost).toDP(10).toString(),
-                totalCost: new Decimal(existingCost.totalCost).plus(cost.totalCost).toDP(10).toString(),
-                tokenMetrics: {
-                    promptTokens: existingCost.tokenMetrics.promptTokens + cost.tokenMetrics.promptTokens,
-                    completionTokens: existingCost.tokenMetrics.completionTokens + cost.tokenMetrics.completionTokens,
-                    totalTokens: existingCost.tokenMetrics.totalTokens + cost.tokenMetrics.totalTokens
-                }
-            };
-            state.set(stateKey, aggregatedCost);
+            try {
+                // Aggregate with existing cost using Decimal
+                const aggregatedCost = {
+                    promptCost: new Decimal(existingCost.promptCost).plus(cost.promptCost).toDP(10).toString(),
+                    completionCost: new Decimal(existingCost.completionCost).plus(cost.completionCost).toDP(10).toString(),
+                    totalCost: new Decimal(existingCost.totalCost).plus(cost.totalCost).toDP(10).toString(),
+                    tokenMetrics: {
+                        promptTokens: existingCost.tokenMetrics.promptTokens + cost.tokenMetrics.promptTokens,
+                        completionTokens: existingCost.tokenMetrics.completionTokens + cost.tokenMetrics.completionTokens,
+                        totalTokens: existingCost.tokenMetrics.totalTokens + cost.tokenMetrics.totalTokens
+                    }
+                };
+                state.set(stateKey, aggregatedCost);
+            }
+            catch (error) {
+                // If aggregation fails, just use the new cost
+                state.set(stateKey, cost);
+            }
         }
         else {
             // First time tracking this agent's cost

package/dist/agents/index.js

@@ -68,8 +68,15 @@ class StatefulAxAgent extends AxAgent {
     // Get the usage cost for the most recent run of the agent
     getLastUsageCost() {
         const { modelUsage, modelInfo, defaults } = this.axai;
-        const currentModelInfo = modelInfo?.find((m) => m.name === defaults.model);
-        if (!currentModelInfo || !modelUsage) {
+        // Check if all required properties exist
+        if (!modelUsage?.promptTokens || !modelUsage?.completionTokens) {
+            return null;
+        }
+        if (!modelInfo || !defaults?.model) {
+            return null;
+        }
+        const currentModelInfo = modelInfo.find((m) => m.name === defaults.model);
+        if (!currentModelInfo?.promptTokenCostPer1M || !currentModelInfo?.completionTokenCostPer1M) {
             return null;
         }
         return StateFulAxAgentUsage.calculateCost(modelUsage, currentModelInfo);

package/dist/types.d.ts

@@ -32,13 +32,22 @@ type FunctionRegistryType = {
 };
 /**
  * The usage metrics of the model.
- * promptTokens: number;
- * completionTokens: number;
+ * Supports both direct token properties and nested tokens structure
  */
-interface ModelUsage {
-    promptTokens: number;
-    completionTokens: number;
+interface ModelUsageBase {
+    promptTokens?: number;
+    completionTokens?: number;
+}
+interface ModelUsageNested {
+    ai?: string;
+    model?: string;
+    tokens?: {
+        totalTokens?: number;
+        promptTokens: number;
+        completionTokens: number;
+    };
 }
+type ModelUsage = ModelUsageBase & ModelUsageNested;
 /**
  * The published cost for using the model.
  * promptTokenCostPer1M: number;

package/examples/streaming.ts

@@ -11,10 +11,10 @@ const config = {
       description: "Completes a user specified task",
       signature:
         'question:string "a question to be answered" -> answer:string "the answer to the question"',
-      provider: "openai", 
-      providerKeyName: "OPENAI_API_KEY",
+      provider: "google-gemini", 
+      providerKeyName: "GEMINI_API_KEY",
       ai: {
-        model: "gpt-4o-mini",
+        model: "gemini-2.0-flash",
         maxTokens: 1000,
         temperature: 0,
       },

package/package.json

@@ -1,7 +1,7 @@
 {
   "type": "module",
   "name": "@amitdeshmukh/ax-crew",
-  "version": "3.7.1",
+  "version": "3.8.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",

package/src/agents/agentUseCosts.ts

@@ -15,52 +15,75 @@ import type {
 export class StateFulAxAgentUsage {
   static STATE_KEY_PREFIX = 'agent_usage_';
 
-  static calculateCost(modelUsage: ModelUsage, modelInfo: ModelInfo): UsageCost {
-    const { promptTokens, completionTokens } = modelUsage;
+  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;
+    const completionTokens = (modelUsage as any).tokens?.completionTokens ?? modelUsage.completionTokens;
     const { promptTokenCostPer1M, completionTokenCostPer1M } = modelInfo;
 
-    // Use Decimal for precise calculations
-    const promptCost = new Decimal(promptTokens)
-      .div(1000000)
-      .mul(promptTokenCostPer1M)
-      .toDP(10);  // Keep 10 decimal places
+    // Return null instead of throwing errors for invalid values
+    if (typeof promptTokens !== 'number' || isNaN(promptTokens) ||
+        typeof completionTokens !== 'number' || isNaN(completionTokens) ||
+        typeof promptTokenCostPer1M !== 'number' || isNaN(promptTokenCostPer1M) ||
+        typeof completionTokenCostPer1M !== 'number' || isNaN(completionTokenCostPer1M)) {
+      return null;
+    }
 
-    const completionCost = new Decimal(completionTokens)
-      .div(1000000)
-      .mul(completionTokenCostPer1M)
-      .toDP(10);
+    try {
+      // Use Decimal for precise calculations
+      const promptCost = new Decimal(promptTokens)
+        .div(1000000)
+        .mul(promptTokenCostPer1M)
+        .toDP(10);  // Keep 10 decimal places
 
-    const totalCost = promptCost.plus(completionCost).toDP(10);
+      const completionCost = new Decimal(completionTokens)
+        .div(1000000)
+        .mul(completionTokenCostPer1M)
+        .toDP(10);
 
-    return {
-      promptCost: promptCost.toString(),
-      completionCost: completionCost.toString(),
-      totalCost: totalCost.toString(),
-      tokenMetrics: {
-        promptTokens,
-        completionTokens,
-        totalTokens: promptTokens + completionTokens
-      }
-    };
+      const totalCost = promptCost.plus(completionCost).toDP(10);
+
+      return {
+        promptCost: promptCost.toString(),
+        completionCost: completionCost.toString(),
+        totalCost: totalCost.toString(),
+        tokenMetrics: {
+          promptTokens,
+          completionTokens,
+          totalTokens: promptTokens + completionTokens
+        }
+      };
+    } catch (error) {
+      // If any decimal calculation fails, return null instead of throwing
+      return null;
+    }
   }
 
-  static trackCostInState(agentName: string, cost: UsageCost, state: StateInstance) {
+  static trackCostInState(agentName: string, cost: UsageCost | null, state: StateInstance) {
+    // If cost is null, skip tracking
+    if (!cost) return;
+
     const stateKey = `${this.STATE_KEY_PREFIX}${agentName}`;
     const existingCost = state.get(stateKey) as UsageCost | undefined;
 
     if (existingCost) {
-      // Aggregate with existing cost using Decimal
-      const aggregatedCost: UsageCost = {
-        promptCost: new Decimal(existingCost.promptCost).plus(cost.promptCost).toDP(10).toString(),
-        completionCost: new Decimal(existingCost.completionCost).plus(cost.completionCost).toDP(10).toString(),
-        totalCost: new Decimal(existingCost.totalCost).plus(cost.totalCost).toDP(10).toString(),
-        tokenMetrics: {
-          promptTokens: existingCost.tokenMetrics.promptTokens + cost.tokenMetrics.promptTokens,
-          completionTokens: existingCost.tokenMetrics.completionTokens + cost.tokenMetrics.completionTokens,
-          totalTokens: existingCost.tokenMetrics.totalTokens + cost.tokenMetrics.totalTokens
-        }
-      };
-      state.set(stateKey, aggregatedCost);
+      try {
+        // Aggregate with existing cost using Decimal
+        const aggregatedCost: UsageCost = {
+          promptCost: new Decimal(existingCost.promptCost).plus(cost.promptCost).toDP(10).toString(),
+          completionCost: new Decimal(existingCost.completionCost).plus(cost.completionCost).toDP(10).toString(),
+          totalCost: new Decimal(existingCost.totalCost).plus(cost.totalCost).toDP(10).toString(),
+          tokenMetrics: {
+            promptTokens: existingCost.tokenMetrics.promptTokens + cost.tokenMetrics.promptTokens,
+            completionTokens: existingCost.tokenMetrics.completionTokens + cost.tokenMetrics.completionTokens,
+            totalTokens: existingCost.tokenMetrics.totalTokens + cost.tokenMetrics.totalTokens
+          }
+        };
+        state.set(stateKey, aggregatedCost);
+      } catch (error) {
+        // If aggregation fails, just use the new cost
+        state.set(stateKey, cost);
+      }
     } else {
       // First time tracking this agent's cost
       state.set(stateKey, cost);

package/src/agents/index.ts

@@ -144,9 +144,19 @@ class StatefulAxAgent extends AxAgent<any, any> {
   // Get the usage cost for the most recent run of the agent
   getLastUsageCost(): UsageCost | null {
     const { modelUsage, modelInfo, defaults } = this.axai;
-    const currentModelInfo = modelInfo?.find((m: { name: string }) => m.name === defaults.model);
     
-    if (!currentModelInfo || !modelUsage) {
+    // Check if all required properties exist
+    if (!modelUsage?.promptTokens || !modelUsage?.completionTokens) {
+      return null;
+    }
+
+    if (!modelInfo || !defaults?.model) {
+      return null;
+    }
+
+    const currentModelInfo = modelInfo.find((m: { name: string }) => m.name === defaults.model);
+    
+    if (!currentModelInfo?.promptTokenCostPer1M || !currentModelInfo?.completionTokenCostPer1M) {
       return null;
     }
 

package/src/types.ts

@@ -36,14 +36,25 @@ type FunctionRegistryType = {
 
 /**
  * The usage metrics of the model.
- * promptTokens: number;
- * completionTokens: number;
+ * Supports both direct token properties and nested tokens structure
  */
-interface ModelUsage {
-  promptTokens: number;
-  completionTokens: number;
+interface ModelUsageBase {
+  promptTokens?: number;
+  completionTokens?: number;
 }
 
+interface ModelUsageNested {
+  ai?: string;
+  model?: string;
+  tokens?: {
+    totalTokens?: number;
+    promptTokens: number;
+    completionTokens: number;
+  };
+}
+
+type ModelUsage = ModelUsageBase & ModelUsageNested;
+
 /**
  * The published cost for using the model.
  * promptTokenCostPer1M: number;