@amitdeshmukh/ax-crew 2.0.3 → 2.0.7

package/CHANGELOG.md

@@ -0,0 +1,23 @@
+# Changelog
+
+This Changelog format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+## [Unreleased]
+
+### Added
+- Change [agent_config.example.yaml](agent_config.example.yaml) to JSON format
+- Add support for providing examples via `setExamples` method
+
+## [2.0.7] - 2024-11-23
+
+### Added
+- Release of version 2.0.7
+- Support for [AxLLM](https://axllm.dev) versions > 10.0.0
+- Changelog
+
+### Changed
+- Updated README.md
+
+### Fixed
+- Updates to [dateTime.ts](src/dateTime.ts) to enable use in Gemini models
+

package/README.md

@@ -12,9 +12,15 @@ This repo simplifies development of [AxLLM](https://axllm.dev) AI Agents by usin
 ## Getting Started
 
 ### Installation
+Install this package:
 ```bash
 npm install @amitdeshmukh/ax-crew
 ```
+AxLLM is a peer dependency, so you will need to install it separately.
+
+```bash
+npm install @ax-llm/ax@latest
+```
 
 ### Environment Setup
 Refer to the [.env.example](.env.example) file for the required environment variables. These will need to be set in the environment where the agents are run.
@@ -139,3 +145,7 @@ const answer = managerResponse.answer;
 console.log(`\n\nPlan: ${plan}`);
 console.log(`\n\nAnswer: ${answer}`);
 ```
+
+## Changelog
+
+See [CHANGELOG.md](CHANGELOG.md) for a list of changes and version updates.

package/agent_config.example.yaml

@@ -3,13 +3,13 @@ crew:
 - name: Planner
   description: Creates a plan to complete a task
   signature: task:string "a task to be completed" -> plan:string "a plan to execute the task in 5 steps or less"
-  provider: anthropic
-  provider_key_name: ANTHROPIC_API_KEY
+  provider: google-gemini
+  provider_key_name: GEMINI_API_KEY
   ai:
-    model: claude-3-5-sonnet-20240620
+    model: gemini-1.5-flash
     temperature: 0
   options:
-    debug: true
+    debug: false
 
 - name: Calculator
   description: Solves math problems
@@ -26,10 +26,10 @@ crew:
 - name: Manager
   description: Answers questions from the user
   signature: question:string "a question from a user", plan:string "a suggested plan to answer the question" -> answer:string "the answer"
-  provider: anthropic
-  provider_key_name: ANTHROPIC_API_KEY
+  provider: openai
+  provider_key_name: OPENAI_API_KEY
   ai:
-    model: claude-3-5-sonnet-20240620
+    model: gpt-4o-mini
     temperature: 0
   options:
     debug: true

package/dist/agents/index.js

@@ -1,16 +1,20 @@
-import { v4 as uuidv4 } from 'uuid';
-import { AxAgent } from '@ax-llm/ax';
-import { getAgentConfigParams } from './agentConfig.js';
-import { createState } from '../state/createState.js';
+import { v4 as uuidv4 } from "uuid";
+import { AxAgent } from "@ax-llm/ax";
+import { getAgentConfigParams } from "./agentConfig.js";
+import { createState } from "../state/index.js";
 // Extend the AxAgent class to include shared state functionality
 class StatefulAxAgent extends AxAgent {
     constructor(ai, options, state) {
         const formattedOptions = {
             ...options,
-            functions: options.functions?.map(fn => typeof fn === 'function' ? fn() : fn)
+            functions: options.functions?.map((fn) => typeof fn === "function" ? fn() : fn),
         };
-        super(ai, formattedOptions);
+        super(formattedOptions);
         this.state = state;
+        this.axai = ai;
+    }
+    async forward(input, options) {
+        return super.forward(this.axai, input, options);
     }
 }
 /**
@@ -49,7 +53,7 @@ class AxCrew {
                     description,
                     signature,
                     functions: functions.filter((fn) => fn !== undefined),
-                    agents: subAgents.filter((agent) => agent !== undefined)
+                    agents: subAgents.filter((agent) => agent !== undefined),
                 }, this.state);
                 return agent;
             }
@@ -81,7 +85,7 @@ class AxCrew {
      * @returns {Map<string, StatefulAxAgent> | null} A map of agent names to their corresponding instances.
      */
     addAgentsToCrew(agentNames) {
-        agentNames.forEach(agentName => {
+        agentNames.forEach((agentName) => {
             this.addAgent(agentName);
         });
         return this.agents;

package/dist/functions/dateTime.js

@@ -3,10 +3,25 @@ export const CurrentDateTime = {
     description: 'Get the current date and time.',
     parameters: {
         type: 'object',
-        properties: {},
-        required: []
+        properties: {
+            format: {
+                type: 'string',
+                description: 'The format of the date and time to return.',
+                enum: ['iso', 'datetime', 'date']
+            }
+        }
     },
-    func: () => new Date().toISOString()
+    func: ({ format }) => {
+        const now = new Date();
+        switch (format) {
+            case 'datetime':
+                return now.toLocaleString();
+            case 'date':
+                return now.toLocaleDateString();
+            default:
+                return now.toISOString();
+        }
+    }
 };
 export const DaysBetweenDates = {
     name: 'DaysBetweenDates',
@@ -16,11 +31,11 @@ export const DaysBetweenDates = {
         properties: {
             startDate: {
                 type: 'string',
-                description: 'The start date in ISO format.'
+                description: 'The start date in ISO format. Must be on or before the end date.'
             },
             endDate: {
                 type: 'string',
-                description: 'The end date in ISO format.'
+                description: 'The end date in ISO format. Must be on or after the start date.'
             }
         },
         required: ['startDate', 'endDate']

package/dist/functions/index.js

@@ -1,6 +1,7 @@
 import { CurrentDateTime, DaysBetweenDates } from './dateTime.js';
+// Built-in functions
 const AxCrewFunctions = {
     CurrentDateTime,
-    DaysBetweenDates,
+    DaysBetweenDates
 };
 export { AxCrewFunctions };

package/package.json

@@ -1,7 +1,7 @@
 {
   "type": "module",
   "name": "@amitdeshmukh/ax-crew",
-  "version": "2.0.3",
+  "version": "2.0.7",
   "description": "Build and launch a crew of AI agents with shared state. Built with axllm.dev",
   "main": "dist/index.js",
   "engines": {
@@ -14,10 +14,11 @@
   "dependencies": {
     "dotenv": "^16.4.5",
     "js-yaml": "^4.1.0",
+    "upgrade": "^1.1.0",
     "uuid": "^10.0.0"
   },
   "peerDependencies": {
-    "@ax-llm/ax": "^9.0.23"
+    "@ax-llm/ax": "^10.0.0"
   },
   "devDependencies": {
     "@types/js-yaml": "^4.0.9",
@@ -25,5 +26,13 @@
     "@types/uuid": "^10.0.0",
     "typescript": "^5.5.3"
   },
+  "homepage": "https://github.com/amitdeshmukh/ax-crew#readme",
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/amitdeshmukh/ax-crew.git"
+  },
+  "bugs": {
+    "url": "https://github.com/amitdeshmukh/ax-crew/issues"
+  },  
   "license": "MIT"
 }

package/src/agents/index.ts

@@ -1,9 +1,14 @@
-import { v4 as uuidv4 } from 'uuid';
-import { AxAgent, AxAI } from '@ax-llm/ax';
-import type { AxSignature, AxAgentic, AxFunction } from '@ax-llm/ax';
-import { getAgentConfigParams } from './agentConfig.js';
-import { FunctionRegistryType } from '../functions/index.js'; 
-import { createState, StateInstance } from '../state/createState.js';
+import { v4 as uuidv4 } from "uuid";
+import { AxAgent, AxAI } from "@ax-llm/ax";
+import type {
+  AxSignature,
+  AxAgentic,
+  AxFunction,
+  AxProgramForwardOptions,
+} from "@ax-llm/ax";
+import { getAgentConfigParams } from "./agentConfig.js";
+import { FunctionRegistryType } from "../functions/index.js";
+import { createState, StateInstance } from "../state/index.js";
 
 // Define the interface for the agent configuration
 interface AgentConfigParams {
@@ -11,21 +16,45 @@ interface AgentConfigParams {
   name: string;
   description: string;
   signature: AxSignature;
-  functions: (AxFunction | (new (state: Record<string, any>) => { toFunction: () => AxFunction; }) | undefined)[];
+  functions: (
+    | AxFunction
+    | (new (state: Record<string, any>) => { toFunction: () => AxFunction })
+    | undefined
+  )[];
   subAgentNames: string[];
 }
 
 // Extend the AxAgent class to include shared state functionality
 class StatefulAxAgent extends AxAgent<any, any> {
   state: StateInstance;
+  axai: any;
 
-  constructor(ai: AxAI, options: Readonly<{ name: string; description: string; signature: string | AxSignature; agents?: AxAgentic[] | undefined; functions?: (AxFunction | (() => AxFunction))[] | undefined; }>, state: StateInstance) {
+  constructor(
+    ai: AxAI,
+    options: Readonly<{
+      name: string;
+      description: string;
+      signature: string | AxSignature;
+      agents?: AxAgentic[] | undefined;
+      functions?: (AxFunction | (() => AxFunction))[] | undefined;
+    }>,
+    state: StateInstance
+  ) {
     const formattedOptions = {
       ...options,
-      functions: options.functions?.map(fn => typeof fn === 'function' ? fn() : fn) as AxFunction[] | undefined
+      functions: options.functions?.map((fn) =>
+        typeof fn === "function" ? fn() : fn
+      ) as AxFunction[] | undefined,
     };
-    super(ai, formattedOptions);
+    super(formattedOptions);
     this.state = state;
+    this.axai = ai;
+  }
+  async forward(
+    input: Record<string, any>,
+    options?: Readonly<AxProgramForwardOptions>
+  ): Promise<Record<string, any>> {
+    return super.forward(this.axai, input, options);
   }
 }
 
@@ -66,39 +95,43 @@ class AxCrew {
   createAgent = (agentName: string): StatefulAxAgent => {
     try {
       const agentConfigParams: AgentConfigParams = getAgentConfigParams(
-        agentName, 
+        agentName,
         this.configFilePath,
         this.functionsRegistry,
         this.state
       );
 
       // Destructure with type assertion
-      const { 
-        ai, 
-        name, 
-        description, 
-        signature, 
-        functions,
-        subAgentNames
-      } = agentConfigParams;
+      const { ai, name, description, signature, functions, subAgentNames } =
+        agentConfigParams;
 
       // Get subagents for the AI agent
       const subAgents = subAgentNames.map((subAgentName: string) => {
         if (!this.agents?.get(subAgentName)) {
-          throw new Error(`Sub-agent '${subAgentName}' does not exist in available agents.`);
+          throw new Error(
+            `Sub-agent '${subAgentName}' does not exist in available agents.`
+          );
         }
         return this.agents?.get(subAgentName);
       });
 
       // Create an instance of StatefulAxAgent
-      const agent = new StatefulAxAgent(ai, {
-        name,
-        description,
-        signature,
-        functions: functions.filter((fn): fn is AxFunction => fn !== undefined),
-        agents: subAgents.filter((agent): agent is StatefulAxAgent => agent !== undefined)
-      }, this.state);
-      
+      const agent = new StatefulAxAgent(
+        ai,
+        {
+          name,
+          description,
+          signature,
+          functions: functions.filter(
+            (fn): fn is AxFunction => fn !== undefined
+          ),
+          agents: subAgents.filter(
+            (agent): agent is StatefulAxAgent => agent !== undefined
+          ),
+        },
+        this.state
+      );
+
       return agent;
     } catch (error) {
       console.error(`Failed to create agent '${agentName}':`, error);
@@ -114,7 +147,7 @@ class AxCrew {
     if (this.agents && !this.agents.has(agentName)) {
       this.agents.set(agentName, this.createAgent(agentName));
     }
-  }  
+  }
 
   /**
    * Sets up agents in the crew by name.
@@ -124,7 +157,7 @@ class AxCrew {
    * @returns {Map<string, StatefulAxAgent> | null} A map of agent names to their corresponding instances.
    */
   addAgentsToCrew(agentNames: string[]): Map<string, StatefulAxAgent> | null {
-    agentNames.forEach(agentName => {
+    agentNames.forEach((agentName) => {
       this.addAgent(agentName);
     });
     return this.agents;
@@ -136,7 +169,7 @@ class AxCrew {
   destroy() {
     this.agents = null;
     this.state.reset();
-  }  
+  }
 }
 
-export { AxCrew };
+export { AxCrew };

package/src/functions/dateTime.ts

@@ -5,10 +5,25 @@ export const CurrentDateTime: AxFunction = {
   description: 'Get the current date and time.',
   parameters: {
     type: 'object',
-    properties: {},
-    required: []
+    properties: {
+      format: {
+        type: 'string',
+        description: 'The format of the date and time to return.',
+        enum: ['iso', 'datetime', 'date']
+      }
+    }
   },
-  func: () => new Date().toISOString()
+  func: ({ format }: { format: string }) => {
+    const now = new Date();
+    switch (format) {
+      case 'datetime':
+        return now.toLocaleString();
+      case 'date':
+        return now.toLocaleDateString();
+      default:
+        return now.toISOString();
+    }
+  }
 }
 
 export const DaysBetweenDates: AxFunction = {
@@ -19,11 +34,11 @@ export const DaysBetweenDates: AxFunction = {
     properties: {
       startDate: { 
         type: 'string',
-        description: 'The start date in ISO format.'
+        description: 'The start date in ISO format. Must be on or before the end date.'
       },
       endDate: {
         type: 'string',
-        description: 'The end date in ISO format.'
+        description: 'The end date in ISO format. Must be on or after the start date.'
       }
     },
     required: ['startDate', 'endDate']

package/src/functions/index.ts

@@ -1,14 +1,15 @@
 import { AxFunction } from '@ax-llm/ax';
 import { CurrentDateTime, DaysBetweenDates } from './dateTime.js';
 
-// FunctionRegistryType
+// FunctionRegistryType to register custom functions
 type FunctionRegistryType = {
   [key: string]: AxFunction | { new(state: Record<string, any>): { toFunction: () => AxFunction } };
 };
 
+// Built-in functions
 const AxCrewFunctions = {
   CurrentDateTime,
-  DaysBetweenDates,
+  DaysBetweenDates
 };
 
 export { AxCrewFunctions, FunctionRegistryType };

package/test.js

@@ -1,21 +1,61 @@
-import { AxCrew, AxCrewFunctions } from '@amitdeshmukh/ax-crew';
+import { AxCrew, AxCrewFunctions } from "./dist/index.js";
 
 // Create a new instance of AxCrew
-const crew = new AxCrew('./agent_config.example.yaml', AxCrewFunctions);
-crew.addAgentsToCrew(['Planner', 'Manager']);
+const crew = new AxCrew("./agent_config.example.yaml", AxCrewFunctions);
+crew.addAgentsToCrew(["Planner", "Calculator", "Manager"]);
 
 // Get agent instances
 const Planner = crew.agents.get("Planner");
 const Manager = crew.agents.get("Manager");
 
+// Calculate costs
+const calculateAgentCost = (agent) => {
+  const modelUsed = agent.aiInstance.models.model;
+  const usedTokens = agent.aiInstance.modelUsage;
+  const modelInfo = agent.aiInstance.modelInfo?.find(
+    (model) => model.name === modelUsed);
+  
+  const promptCost =
+    (usedTokens.promptTokens / 1000000) * modelInfo?.promptTokenCostPer1M;
+  const completionCost =
+    (usedTokens.completionTokens / 1000000) *
+    modelInfo?.completionTokenCostPer1M;
+  const totalCost = promptCost + completionCost;
+
+  return {
+    totalCost,
+    promptTokens: usedTokens.promptTokens,
+    completionTokens: usedTokens.completionTokens,
+  };
+};
+
 // User query
 const userQuery = "whats the square root of the number of days between now and Christmas";
-
 console.log(`\n\nQuestion: ${userQuery}`);
 
 // Forward the user query to the agents
-const planResponse = await Planner.forward({ task: userQuery });
-const managerResponse = await Manager.forward({ question: userQuery, plan: planResponse.plan });
+const planResponse = await Planner.forward({
+  task: userQuery
+});
+
+const planner = calculateAgentCost(Planner);
+console.log('\nPlanner Usage:');
+console.log(`Cost: $${planner.totalCost.toFixed(6)}`);
+console.log(`Prompt Tokens: ${planner.promptTokens}`);
+console.log(`Completion Tokens: ${planner.completionTokens}`);
+console.log(`Total Tokens: ${planner.promptTokens + planner.completionTokens}`);
+
+const managerResponse = await Manager.forward({
+  question: userQuery,
+  plan: planResponse.plan,
+});
+
+const manager = calculateAgentCost(Manager);
+console.log('\nManager Usage:');
+console.log(`Cost: $${manager.totalCost.toFixed(4)}`);
+console.log(`Prompt Tokens: ${manager.promptTokens}`);
+console.log(`Completion Tokens: ${manager.completionTokens}`);
+console.log(`Total Tokens: ${manager.promptTokens + manager.completionTokens}`);
 
 // Get and print the plan and answer from the agents
 const plan = planResponse.plan;