@amitdeshmukh/ax-crew 3.3.4 → 3.5.1

package/CHANGELOG.md

@@ -5,11 +5,23 @@ This Changelog format is based on [Keep a Changelog]
 adheres to [Semantic Versioning](https://semver.org/spec/
 v2.0.0.html).
 
-## [3.3.4] - 2024-12-31
+## [3.5.2] - 2025-02-03
 
-### Removed
-- Removed unused `js-yaml` dependency
-- Updated `README.md` to highlight AxLLM 10.0.9 as peer dependency
+### Added
+- Enhanced cost tracking with precise decimal calculations for agent and crew usage
+- Improved token metrics aggregation across multiple agent runs
+- Added support for per-agent and total crew cost analysis
+
+### Fixed
+- Added missing decimal.js dependency installation to resolve type declarations error
+
+## [3.3.4] - 2025-02-03
+
+### Changed
+- Fixed sub-agent function calling in StatefulAxAgent to properly handle AI configuration
+- Improved TypeScript function overloads for the forward method to provide better type safety and readability
+- Removed redundant AI parameter passing in sub-agent calls
+- Added an [example](examples/write-post-and-publish-to-wordpress.ts) of using AxCrew to write a post and publish it to WordPress
 
 ## [3.3.3] - 2024-12-12
 

package/README.md

@@ -16,12 +16,10 @@ Install this package:
 ```bash
 npm install @amitdeshmukh/ax-crew
 ```
-AxLLM is a peer dependency, so you will need to install it separately. 
-
-**Note:** Currently, we support AxLLM v10.0.9. We are working on updating this package to support the latest version of AxLLM as it introduces some breaking changes.
+AxLLM is a peer dependency, so you will need to install it separately.
 
 ```bash
-npm install @ax-llm/ax@10.0.9
+npm install @ax-llm/ax@latest
 ```
 
 ### TypeScript Support
@@ -76,17 +74,22 @@ const myFunctions: FunctionRegistryType = {
 // Create crew with type checking
 const crew = new AxCrew(config, myFunctions);
 
-// Type-safe state management
+// Set and get state
 crew.state.set('key', 'value');
 const value: string = crew.state.get('key');
 
-// Type-safe agent management
+// Add agents to the crew
 const agents = crew.addAgentsToCrew(['Planner']);
 const planner = agents?.get('Planner');
 
 if (planner) {
-  // Type-safe agent usage
+  // Agent usage with function overloads
+  // Direct usage - AI config from agent construction is used
   const response = await planner.forward({ task: "Plan something" });
+  
+  // Sub-agent usage - when used by another agent (AI is ignored and agent's own config is used)
+  const subAgentResponse = await planner.forward(ai, { task: "Plan something" });
+  
   const cost = planner.getUsageCost();
   
   if (cost) {
@@ -313,19 +316,20 @@ console.log(`\n\nAnswer: ${answer}`);
 
 ### Tracking Usage Costs
 
-You can track the usage costs for each agent's execution using the `getUsageCost` method. This is particularly useful for monitoring API costs and usage patterns.
+The package provides precise cost tracking capabilities for monitoring API usage across individual agents and the entire crew. Costs are calculated using high-precision decimal arithmetic to ensure accuracy.
 
 ```javascript
 // After running an agent's forward method
 const response = await Planner.forward({ task: userQuery });
-const cost = Planner.getUsageCost();
 
-console.log(cost);
+// Get individual agent costs
+const agentCost = Planner.getUsageCost();
+console.log(agentCost);
 /* Output example:
 {
-  promptCost: 0.00036375,
-  completionCost: 0.00061,
-  totalCost: 0.00097375,
+  promptCost: "0.0003637500000",
+  completionCost: "0.0006100000000",
+  totalCost: "0.0009737500000",
   tokenMetrics: {
     promptTokens: 291,
     completionTokens: 122,
@@ -333,8 +337,40 @@ console.log(cost);
   }
 }
 */
+
+// Get aggregated costs for all agents in the crew
+const crewCosts = crew.getAggregatedCosts();
+console.log(crewCosts);
+/* Output example:
+{
+  totalCost: "0.0025482500000",
+  byAgent: {
+    "Planner": { ... },
+    "Calculator": { ... },
+    "Manager": { ... }
+  },
+  aggregatedMetrics: {
+    promptTokens: 850,
+    completionTokens: 324,
+    totalTokens: 1174,
+    promptCost: "0.0010625000000",
+    completionCost: "0.0014857500000"
+  }
+}
+*/
+
+// Reset cost tracking if needed
+crew.resetCosts();
 ```
 
+Cost tracking features:
+- High-precision decimal calculations using decimal.js
+- Per-agent cost breakdown
+- Aggregated crew-wide metrics
+- Token usage statistics
+- Support for different pricing tiers per model
+- Persistent cost tracking across multiple agent runs
+
 ## Changelog
 
 See [CHANGELOG.md](CHANGELOG.md) for a list of changes and version updates.

package/dist/agents/agentUseCosts.d.ts

@@ -1,3 +1,4 @@
+import type { StateInstance } from '../state/index.js';
 export interface ModelUsage {
     promptTokens: number;
     completionTokens: number;
@@ -7,9 +8,9 @@ export interface ModelInfo {
     completionTokenCostPer1M: number;
 }
 export interface UsageCost {
-    promptCost: number;
-    completionCost: number;
-    totalCost: number;
+    promptCost: string;
+    completionCost: string;
+    totalCost: string;
     tokenMetrics: {
         promptTokens: number;
         completionTokens: number;
@@ -17,5 +18,19 @@ export interface UsageCost {
     };
 }
 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 getAggregatedCosts(state: StateInstance): {
+        totalCost: string;
+        byAgent: Record<string, UsageCost>;
+        aggregatedMetrics: {
+            promptTokens: number;
+            completionTokens: number;
+            totalTokens: number;
+            promptCost: string;
+            completionCost: string;
+        };
+    };
+    static resetCosts(state: StateInstance): void;
 }

package/dist/agents/agentUseCosts.js

@@ -1,13 +1,23 @@
+import { Decimal } from 'decimal.js';
 // Utility class to handle usage related functionality
 export class StateFulAxAgentUsage {
     static calculateCost(modelUsage, modelInfo) {
         const { promptTokens, completionTokens } = modelUsage;
         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: (promptTokens / 1000000) * promptTokenCostPer1M,
-            completionCost: (completionTokens / 1000000) * completionTokenCostPer1M,
-            totalCost: ((promptTokens / 1000000) * promptTokenCostPer1M) +
-                ((completionTokens / 1000000) * completionTokenCostPer1M),
+            promptCost: promptCost.toString(),
+            completionCost: completionCost.toString(),
+            totalCost: totalCost.toString(),
             tokenMetrics: {
                 promptTokens,
                 completionTokens,
@@ -15,4 +25,68 @@ export class StateFulAxAgentUsage {
             }
         };
     }
+    static trackCostInState(agentName, cost, state) {
+        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);
+        }
+        else {
+            // First time tracking this agent's cost
+            state.set(stateKey, cost);
+        }
+    }
+    static getAggregatedCosts(state) {
+        const allState = state.getAll();
+        const agentCosts = {};
+        let totalPromptTokens = 0;
+        let totalCompletionTokens = 0;
+        let totalPromptCost = new Decimal(0);
+        let totalCompletionCost = new Decimal(0);
+        // Aggregate costs from all agents
+        Object.entries(allState).forEach(([key, value]) => {
+            if (key.startsWith(this.STATE_KEY_PREFIX)) {
+                const agentName = key.replace(this.STATE_KEY_PREFIX, '');
+                const cost = value;
+                agentCosts[agentName] = cost;
+                totalPromptTokens += cost.tokenMetrics.promptTokens;
+                totalCompletionTokens += cost.tokenMetrics.completionTokens;
+                totalPromptCost = totalPromptCost.plus(cost.promptCost);
+                totalCompletionCost = totalCompletionCost.plus(cost.completionCost);
+            }
+        });
+        // Calculate total cost from the sum of prompt and completion costs
+        const totalCost = totalPromptCost.plus(totalCompletionCost).toDP(10);
+        return {
+            totalCost: totalCost.toString(),
+            byAgent: agentCosts,
+            aggregatedMetrics: {
+                promptTokens: totalPromptTokens,
+                completionTokens: totalCompletionTokens,
+                totalTokens: totalPromptTokens + totalCompletionTokens,
+                promptCost: totalPromptCost.toDP(10).toString(),
+                completionCost: totalCompletionCost.toDP(10).toString()
+            }
+        };
+    }
+    static resetCosts(state) {
+        const allState = state.getAll();
+        Object.keys(allState).forEach(key => {
+            if (key.startsWith(this.STATE_KEY_PREFIX)) {
+                state.set(key, undefined);
+            }
+        });
+    }
 }
+StateFulAxAgentUsage.STATE_KEY_PREFIX = 'agent_usage_';

package/dist/agents/index.d.ts

@@ -3,10 +3,11 @@ import type { AxSignature, AxAgentic, AxFunction, AxProgramForwardOptions } from
 import type { AgentConfigInput } from "./agentConfig.js";
 import { FunctionRegistryType } from "../functions/index.js";
 import { StateInstance } from "../state/index.js";
-import { UsageCost } from "./agentUseCosts.js";
+import { StateFulAxAgentUsage, UsageCost } from "./agentUseCosts.js";
 declare class StatefulAxAgent extends AxAgent<any, any> {
     state: StateInstance;
     axai: any;
+    private agentName;
     constructor(ai: AxAI, options: Readonly<{
         name: string;
         description: string;
@@ -15,8 +16,11 @@ declare class StatefulAxAgent extends AxAgent<any, any> {
         functions?: (AxFunction | (() => AxFunction))[] | undefined;
         examples?: Array<Record<string, any>> | undefined;
     }>, state: StateInstance);
-    forward(input: Record<string, any>, options?: Readonly<AxProgramForwardOptions>): Promise<Record<string, any>>;
+    forward(values: Record<string, any>, options?: Readonly<AxProgramForwardOptions>): Promise<Record<string, any>>;
+    forward(ai: AxAI, values: Record<string, any>, options?: Readonly<AxProgramForwardOptions>): Promise<Record<string, any>>;
     getUsageCost(): UsageCost | null;
+    getAggregatedCosts(): ReturnType<typeof StateFulAxAgentUsage.getAggregatedCosts>;
+    resetCosts(): void;
 }
 /**
  * Represents a crew of agents with shared state functionality.

package/dist/agents/index.js

@@ -14,14 +14,29 @@ class StatefulAxAgent extends AxAgent {
         super(formattedOptions);
         this.state = state;
         this.axai = ai;
+        this.agentName = options.name;
         // Set examples if provided
         if (examples && examples.length > 0) {
-            super.setExamples(examples);
+            this.setExamples(examples);
         }
     }
-    // Override the forward method
-    async forward(input, options) {
-        return super.forward(this.axai, input, options);
+    // Implementation
+    async forward(first, second, third) {
+        // Sub-agent case (called with AI service)
+        let result;
+        if ('apiURL' in first) {
+            result = await super.forward(this.axai, second, third);
+        }
+        else {
+            // Direct call case
+            result = await super.forward(this.axai, first, second);
+        }
+        // Track costs in state after the call
+        const cost = this.getUsageCost();
+        if (cost) {
+            StateFulAxAgentUsage.trackCostInState(this.agentName, cost, this.state);
+        }
+        return result;
     }
     // Get the usage cost for a run of the agent
     getUsageCost() {
@@ -32,6 +47,14 @@ class StatefulAxAgent extends AxAgent {
         }
         return StateFulAxAgentUsage.calculateCost(modelUsage, currentModelInfo);
     }
+    // Get aggregated costs for all agents in the crew
+    getAggregatedCosts() {
+        return StateFulAxAgentUsage.getAggregatedCosts(this.state);
+    }
+    // Reset all cost tracking
+    resetCosts() {
+        StateFulAxAgentUsage.resetCosts(this.state);
+    }
 }
 /**
  * Represents a crew of agents with shared state functionality.

package/examples/README.md

@@ -0,0 +1,162 @@
+# AxCrew Examples
+
+This directory contains example implementations showing how to use AxCrew in different scenarios. Each example demonstrates specific features and use cases.
+
+## Prerequisites
+
+Before running any example, make sure you have:
+
+1. Node.js >= 21.0.0 installed
+2. Required API keys set in your `.env` file:
+   ```env
+   ANTHROPIC_API_KEY=your_anthropic_key
+   OPENAI_API_KEY=your_openai_key
+   GEMINI_API_KEY=your_gemini_key
+   # Add other provider keys as needed
+   ```
+
+## Additional Dependencies
+
+Some examples require additional npm modules. Install them based on which examples you want to run:
+
+```bash
+# For WordPress example
+npm install node-fetch https
+```
+
+## Examples
+
+### 1. Basic Researcher-Writer
+[`basic-researcher-writer.ts`](./basic-researcher-writer.ts)
+
+Demonstrates how to create a simple crew with two agents:
+- A researcher agent that finds information
+- A writer agent that creates content based on the research
+
+**Required Dependencies:** None beyond core package
+
+```typescript
+import { AxCrew } from "@amitdeshmukh/ax-crew";
+
+const crew = new AxCrew({
+  crew: [
+    {
+      name: "researcher",
+      // ... configuration
+    },
+    {
+      name: "writer",
+      // ... configuration
+    }
+  ]
+});
+
+// Usage example in the file
+```
+
+### 2. Solve Math Problem
+[`solve-math-problem.ts`](./solve-math-problem.ts)
+
+Shows how to use the math-solving capabilities with step-by-step explanations:
+- Uses Gemini Pro for code execution
+- Demonstrates how to handle complex mathematical queries
+- Shows usage cost tracking
+
+**Required Dependencies:** 
+```bash
+npm install mathjs
+```
+
+```typescript
+const userQuery = "who is considered as the father of the iphone and what is the 7th root of their year of birth?";
+// See file for implementation
+```
+
+### 3. Write and Publish to WordPress
+[`write-post-and-publish-to-wordpress.ts`](./write-post-and-publish-to-wordpress.ts)
+
+Demonstrates how to:
+- Create content using AI
+- Format it for WordPress
+- Publish directly to a WordPress site
+- Handle API interactions
+
+**Required Dependencies:**
+```bash
+npm install @wordpress/api-fetch @wordpress/url
+```
+
+**Additional Configuration:**
+```env
+WORDPRESS_URL=your_wordpress_site_url
+WORDPRESS_USERNAME=your_username
+WORDPRESS_APP_PASSWORD=your_application_password
+```
+
+### 4. AxCrew Basic Usage
+[`axcrew.ts`](./axcrew.ts)
+
+A basic example showing core AxCrew features:
+- Setting up a crew
+- Adding agents
+- Using shared state
+- Cost tracking
+- Error handling
+
+**Required Dependencies:** None beyond core package
+
+## Running the Examples
+
+1. Clone the repository:
+   ```bash
+   git clone https://github.com/amitdeshmukh/ax-crew.git
+   ```
+
+2. Install dependencies:
+   ```bash
+   cd ax-crew
+   npm install
+   ```
+
+3. Set up your environment variables in `.env`
+
+4. Run an example:
+   ```bash
+   # Using ts-node
+   npx ts-node examples/basic-researcher-writer.ts
+
+   # Or after building
+   npm run build
+   node dist/examples/basic-researcher-writer.js
+   ```
+
+## Cost Tracking
+
+All examples include cost tracking. You can monitor API usage costs:
+
+```typescript
+// After running an agent
+const cost = agent.getUsageCost();
+console.log("Usage cost:", cost);
+
+// Get aggregated costs for all agents
+const totalCosts = agent.getAggregatedCosts();
+console.log("Total costs:", totalCosts);
+```
+
+## Best Practices
+
+1. Always check for and handle errors appropriately
+2. Monitor costs using the built-in cost tracking
+3. Use appropriate model temperatures based on your use case
+4. Leverage shared state for better agent coordination
+5. Consider rate limits of the AI providers
+
+## Contributing
+
+Feel free to contribute more examples by submitting a pull request. Make sure to:
+1. Include clear documentation
+2. Follow the existing code style
+3. Add appropriate error handling
+4. Include cost tracking
+5. Add your example to this README

package/examples/basic-researcher-writer.ts

@@ -0,0 +1,79 @@
+;import { AxCrew } from "../dist/index.js";
+
+// Example agent configuration
+const agentConfig = {
+  crew: [
+    {
+      name: "researcher",
+      description: "A research agent that finds information",
+      signature: "query:string -> research:string",
+      provider: "anthropic",
+      providerKeyName: "ANTHROPIC_API_KEY",
+      ai: {
+        model: "claude-3-opus-20240229"
+      }
+    },
+    {
+      name: "writer",
+      description: "A writing agent that creates content",
+      signature: "topic:string, research:string -> article:string",
+      provider: "anthropic",
+      providerKeyName: "ANTHROPIC_API_KEY",
+      ai: {
+        model: "claude-3-sonnet-20240229"
+      },
+      agents: ["researcher"] // This makes writer use researcher as a sub-agent
+    }
+  ]
+};
+
+// Example usage
+async function main() {
+  // Initialize the crew
+  const crew = new AxCrew(agentConfig);
+  
+  // Add agents to the crew
+  crew.addAgentsToCrew(["researcher", "writer"]);
+  
+  // Get references to our agents
+  const writer = crew.agents?.get("writer");
+  const researcher = crew.agents?.get("researcher");
+  
+  if (!writer || !researcher) {
+    throw new Error("Failed to initialize agents");
+  }
+
+  try {
+    // Use the researcher agent
+    const { research } = await researcher.forward({
+      query: "What are the key benefits of quantum computing?"
+    });
+    
+    // Use the writer agent (which will internally use the researcher)
+    const article = await writer.forward({
+      topic: "Quantum Computing Benefits",
+      research: research
+    });
+
+    // Print the article
+    console.log("Article:", article);
+
+    // Get costs for individual agents
+    const researcherCost = researcher.getUsageCost();
+    console.log("Researcher cost:", researcherCost);
+    
+    // Get aggregated costs for all agents (includes both direct calls and sub-agent calls)
+    const totalCosts = writer.getAggregatedCosts();
+    console.log("Total cost breakdown:", JSON.stringify(totalCosts, null, 2));
+
+    // If you want to start fresh with cost tracking
+    writer.resetCosts();
+    
+  } finally {
+    // Clean up
+    crew.destroy();
+  }
+}
+
+// Run the example
+main().catch(console.error);

package/examples/solve-math-problem.ts

@@ -0,0 +1,74 @@
+import { AxCrew } from "../dist/index.js";
+
+// Define the crew configuration
+const config = {
+  crew: [
+    {
+      name: "MathAgent",
+      description: "Solves math problems",
+      signature:
+        'mathProblem:string "a sentence describing a math problem to be solved using Python code" -> solution:string "the answer to the math problem"',
+      provider: "google-gemini",
+      providerKeyName: "GEMINI_API_KEY",
+      ai: {
+        model: "gemini-1.5-pro",
+        temperature: 0,
+      },
+      options: {
+        debug: false,
+        codeExecution: true,
+      },
+    },
+    {
+      name: "ManagerAgent",
+      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",
+      ai: {
+        model: "gpt-4o-mini",
+        maxTokens: 1000,
+        temperature: 0,
+      },
+      options: {
+        debug: true,
+      },
+      agents: ["MathAgent"]
+    },
+  ],
+};
+
+// Create a new instance of AxCrew with the config
+const crew = new AxCrew(config);
+
+// Add the agents to the crew
+const agents = crew.addAgentsToCrew(["MathAgent", "ManagerAgent"]);
+
+// Get agent instances
+const managerAgent = agents?.get("ManagerAgent");
+
+const userQuery: string =
+  "who is considered as the father of the iphone and what is the 7th root of their year of birth (precision to minimum 5 decimal places)";
+console.log(`\n\nQuestion: ${userQuery}`);
+
+const main = async (): Promise<void> => {
+  if (managerAgent) {
+    // Try with manager agent first
+    const managerResponse = await managerAgent.forward({
+      question: userQuery,
+    });
+
+    console.log(
+      `\nManager's Answer: ${JSON.stringify(managerResponse.answer, null, 2)}\n`
+    );
+
+    // Print usage costs
+    const managerCost = managerAgent.getUsageCost();
+
+    console.log("\nUsage Costs:");
+    console.log("Manager Agent:", managerCost);
+  }
+};
+
+main().catch(console.error);