@amitdeshmukh/ax-crew 3.3.0 → 3.3.2

package/CHANGELOG.md

@@ -5,7 +5,13 @@ This Changelog format is based on [Keep a Changelog]
 adheres to [Semantic Versioning](https://semver.org/spec/
 v2.0.0.html).
 
-## [3.3.0] - 2024-12-11
+## [3.3.1] - 2024-12-11
+
+### Changed
+- Fixed Websearch agent configuration in `agentConfig.json`
+- Refactored agent configuration structure in `src/agents/agentConfig.ts` for better type safety and maintainability
+
+## [3.3.0] - 2024-12-10
 
 ### Added
 - Added `getUsageCost` method to track API usage costs and token metrics

package/agentConfig.json

@@ -62,7 +62,11 @@
       },
       "options": {
         "debug": true,
-        "googleSearchRetrieval": true
+        "googleSearchRetrieval": {
+          "dynamic_retrieval_config": {
+            "mode": "MODE_UNSPECIFIED"
+          }
+        }
       },
       "functions": ["CurrentDateTime", "DaysBetweenDates"]
     },

package/dist/agents/agentConfig.d.ts

@@ -0,0 +1,47 @@
+import type { AxModelConfig, AxFunction, AxSignature } from '@ax-llm/ax';
+import { FunctionRegistryType } from '../functions/index.js';
+type ExtendedAxModelConfig = AxModelConfig & {
+    model: string;
+};
+interface AgentConfig {
+    name: string;
+    description: string;
+    signature: AxSignature;
+    provider: string;
+    providerKeyName?: string;
+    ai: ExtendedAxModelConfig;
+    debug?: boolean;
+    apiURL?: string;
+    options?: Record<string, any>;
+    functions?: string[];
+    agents?: string[];
+    examples?: Array<Record<string, any>>;
+}
+type AgentConfigInput = string | {
+    crew: AgentConfig[];
+};
+/**
+ * Initializes the AI agent using the specified agent name and configuration.
+ * This function parses the agent's configuration, validates the presence of the necessary API key,
+ * and creates an instance of the AI agent with the appropriate settings.
+ *
+ * @param {string} agentName - The identifier for the AI agent to be initialized.
+ * @param {AgentConfigInput} agentConfig - Either a file path to the JSON configuration or a JSON object with crew configuration.
+ * @param {FunctionRegistryType} functions - The functions available to the agent.
+ * @param {Object} state - The state object for the agent.
+ * @returns {Object} An object containing the Agents AI instance, its name, description, signature, functions and subAgentList.
+ * @throws {Error} Throws an error if the agent configuration is missing, the provider is unsupported,
+ * the API key is not found, or the provider key name is not specified in the configuration.
+ */
+declare const getAgentConfigParams: (agentName: string, agentConfig: AgentConfigInput, functions: FunctionRegistryType, state: Record<string, any>) => {
+    ai: any;
+    name: string;
+    description: string;
+    signature: AxSignature;
+    functions: (AxFunction | (new (state: Record<string, any>) => {
+        toFunction: () => AxFunction;
+    }) | undefined)[];
+    subAgentNames: string[];
+    examples: Record<string, any>[];
+};
+export { getAgentConfigParams, AgentConfigInput };

package/dist/agents/agentConfig.js

@@ -133,20 +133,19 @@ const getAgentConfigParams = (agentName, agentConfig, functions, state) => {
         else {
             throw new Error(`Provider key name is missing in the agent configuration`);
         }
-        // Create an instance of the AI agent
+        // Create an instance of the AI agent and set options
         const ai = new AIConstructor({
             apiKey,
             config: agentConfigData.ai,
             options: {
-                debug: agentConfigData.debug || false
+                debug: agentConfigData.debug || false,
+                ...agentConfigData.options
             }
         });
         // If an apiURL is provided in the agent config, set it in the AI agent
         if (agentConfigData.apiURL) {
             ai.setAPIURL(agentConfigData.apiURL);
         }
-        // Set all options from the agent configuration
-        ai.setOptions({ ...agentConfigData.options });
         // Prepare functions for the AI agent
         const agentFunctions = (agentConfigData.functions || [])
             .map(funcName => {

package/dist/agents/agentUseCosts.d.ts

@@ -0,0 +1,21 @@
+export interface ModelUsage {
+    promptTokens: number;
+    completionTokens: number;
+}
+export interface ModelInfo {
+    promptTokenCostPer1M: number;
+    completionTokenCostPer1M: number;
+}
+export interface UsageCost {
+    promptCost: number;
+    completionCost: number;
+    totalCost: number;
+    tokenMetrics: {
+        promptTokens: number;
+        completionTokens: number;
+        totalTokens: number;
+    };
+}
+export declare class StateFulAxAgentUsage {
+    static calculateCost(modelUsage: ModelUsage, modelInfo: ModelInfo): UsageCost;
+}

package/dist/agents/index.d.ts

@@ -0,0 +1,62 @@
+import { AxAgent, AxAI } from "@ax-llm/ax";
+import type { AxSignature, AxAgentic, AxFunction, AxProgramForwardOptions } from "@ax-llm/ax";
+import type { AgentConfigInput } from "./agentConfig.js";
+import { FunctionRegistryType } from "../functions/index.js";
+import { StateInstance } from "../state/index.js";
+import { UsageCost } from "./agentUseCosts.js";
+declare 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;
+        examples?: Array<Record<string, any>> | undefined;
+    }>, state: StateInstance);
+    forward(input: Record<string, any>, options?: Readonly<AxProgramForwardOptions>): Promise<Record<string, any>>;
+    getUsageCost(): UsageCost | null;
+}
+/**
+ * Represents a crew of agents with shared state functionality.
+ */
+declare class AxCrew {
+    private agentConfig;
+    functionsRegistry: FunctionRegistryType;
+    crewId: string;
+    agents: Map<string, StatefulAxAgent> | null;
+    state: StateInstance;
+    /**
+     * Creates an instance of AxCrew.
+     * @param {AgentConfigInput} agentConfig - Either a path to the agent config file or a JSON object with crew configuration.
+     * @param {FunctionRegistryType} [functionsRegistry={}] - The registry of functions to use in the crew.
+     * @param {string} [crewId=uuidv4()] - The unique identifier for the crew.
+     */
+    constructor(agentConfig: AgentConfigInput, functionsRegistry?: FunctionRegistryType, crewId?: string);
+    /**
+     * Factory function for creating an agent.
+     * @param {string} agentName - The name of the agent to create.
+     * @returns {StatefulAxAgent} The created StatefulAxAgent instance.
+     * @throws Will throw an error if the agent creation fails.
+     */
+    createAgent: (agentName: string) => StatefulAxAgent;
+    /**
+     * Adds an agent to the crew by name.
+     * @param {string} agentName - The name of the agent to add.
+     */
+    addAgent(agentName: string): void;
+    /**
+     * Sets up agents in the crew by name.
+     * For an array of Agent names provided, it adds
+     * the agent to the crew if not already present.
+     * @param {string[]} agentNames - An array of agent names to configure.
+     * @returns {Map<string, StatefulAxAgent> | null} A map of agent names to their corresponding instances.
+     */
+    addAgentsToCrew(agentNames: string[]): Map<string, StatefulAxAgent> | null;
+    /**
+     * Cleans up the crew by dereferencing agents and resetting the state.
+     */
+    destroy(): void;
+}
+export { AxCrew };

package/dist/config/index.d.ts

@@ -0,0 +1,5 @@
+interface ProviderApiKeys {
+    [key: string]: string | undefined;
+}
+declare const PROVIDER_API_KEYS: ProviderApiKeys;
+export { PROVIDER_API_KEYS, };

package/dist/functions/dateTime.d.ts

@@ -0,0 +1,3 @@
+import type { AxFunction } from "@ax-llm/ax";
+export declare const CurrentDateTime: AxFunction;
+export declare const DaysBetweenDates: AxFunction;

package/dist/functions/index.d.ts

@@ -0,0 +1,13 @@
+import { AxFunction } from '@ax-llm/ax';
+type FunctionRegistryType = {
+    [key: string]: AxFunction | {
+        new (state: Record<string, any>): {
+            toFunction: () => AxFunction;
+        };
+    };
+};
+declare const AxCrewFunctions: {
+    CurrentDateTime: AxFunction;
+    DaysBetweenDates: AxFunction;
+};
+export { AxCrewFunctions, FunctionRegistryType };

package/dist/index.d.ts

@@ -0,0 +1,3 @@
+import { AxCrew } from './agents/index.js';
+import { AxCrewFunctions, FunctionRegistryType } from './functions/index.js';
+export { AxCrew, AxCrewFunctions, FunctionRegistryType };

package/dist/state/index.d.ts

@@ -0,0 +1,12 @@
+interface StateInstance {
+    set: (key: string, value: any) => void;
+    get: (key: string) => any;
+    getAll: () => {
+        [key: string]: any;
+    };
+    getId: () => string;
+    reset: () => void;
+}
+declare const createState: (id: string) => StateInstance;
+declare const getState: (id: string) => StateInstance | undefined;
+export { createState, getState, StateInstance };

package/index.d.ts

@@ -0,0 +1,7 @@
+export class AxCrew {
+  constructor(configPath: string);
+  addAgent(agent: any): void;
+  createAgent(agent: any): void;
+  addAgentsToCrew(agents: any[]): void;
+  destroy(): void;
+}

package/package.json

@@ -1,9 +1,10 @@
 {
   "type": "module",
   "name": "@amitdeshmukh/ax-crew",
-  "version": "3.3.0",
+  "version": "3.3.2",
   "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",
   "engines": {
     "node": ">=21.0.0"
   },

package/src/agents/agentConfig.ts

@@ -173,21 +173,19 @@ const getAgentConfigParams = (
       throw new Error(`Provider key name is missing in the agent configuration`);
     }
 
-    // Create an instance of the AI agent
+    // Create an instance of the AI agent and set options
     const ai = new AIConstructor({
       apiKey,
       config: agentConfigData.ai,
       options: {
-        debug: agentConfigData.debug || false
+        debug: agentConfigData.debug || false,
+        ...agentConfigData.options
       }
     });
     // If an apiURL is provided in the agent config, set it in the AI agent
     if (agentConfigData.apiURL) {
       ai.setAPIURL(agentConfigData.apiURL);
     }
-
-    // Set all options from the agent configuration
-    ai.setOptions({ ...agentConfigData.options });
     
     // Prepare functions for the AI agent
     const agentFunctions = (agentConfigData.functions || [])

package/src/agents/agentUseCosts.ts

@@ -1,14 +1,14 @@
-interface ModelUsage {
+export interface ModelUsage {
   promptTokens: number;
   completionTokens: number;
 }
 
-interface ModelInfo {
+export interface ModelInfo {
   promptTokenCostPer1M: number;
   completionTokenCostPer1M: number;
 }
 
-interface UsageCost {
+export interface UsageCost {
   promptCost: number;
   completionCost: number;
   totalCost: number;

package/src/agents/index.ts

@@ -10,7 +10,7 @@ import { getAgentConfigParams } from "./agentConfig.js";
 import type { AgentConfigInput } from "./agentConfig.js";
 import { FunctionRegistryType } from "../functions/index.js";
 import { createState, StateInstance } from "../state/index.js";
-import { StateFulAxAgentUsage } from "./agentUseCosts.js";
+import { StateFulAxAgentUsage, UsageCost } from "./agentUseCosts.js";
 
 // Define the interface for the agent configuration
 interface AgentConfigParams {
@@ -70,7 +70,7 @@ class StatefulAxAgent extends AxAgent<any, any> {
   }
 
   // Get the usage cost for a run of the agent
-  getUsageCost() {
+  getUsageCost(): UsageCost | null {
     const { modelUsage, modelInfo, models } = this.axai;
     const currentModelInfo = modelInfo?.find((m: { name: string }) => m.name === models.model);
     

package/test1.js

@@ -6,65 +6,73 @@ const config = {
     {
       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\"",
+      signature:
+        'task:string "a task to be completed" -> plan:string "a plan to execute the task in 5 steps or less"',
       provider: "google-gemini",
       providerKeyName: "GEMINI_API_KEY",
       ai: {
         model: "gemini-1.5-flash",
-        temperature: 0
+        temperature: 0,
       },
       options: {
-        debug: false
-      }
+        debug: false,
+      },
     },
     {
       name: "Calculator",
       description: "Solves math problems",
-      signature: "mathProblem:string \"a math problem to be solved using Python code\" -> solution:string \"the solution to the math problem\"",
+      signature:
+        'mathProblem:string "a math problem to be solved using Python code" -> solution:string "the solution to the math problem"',
       provider: "google-gemini",
       providerKeyName: "GEMINI_API_KEY",
       ai: {
         model: "gemini-1.5-pro",
-        temperature: 0
+        temperature: 0,
       },
       options: {
         debug: true,
-        codeExecution: true
+        codeExecution: true,
       },
-      functions: ["CurrentDateTime", "DaysBetweenDates"]
+      functions: ["CurrentDateTime", "DaysBetweenDates"],
     },
     {
       name: "WebSearch",
-      description: "Searches the web for the latest information using Google search",
-      signature: "webSearchQuery:string \"a query for Google search\" -> webSearchResponse:string \"the result of the search\"",
+      description:
+        "Searches the web for the latest information using Google search",
+      signature:
+        'webSearchQuery:string "a query for Google search" -> webSearchResponse:string "the result of the search"',
       provider: "google-gemini",
       providerKeyName: "GEMINI_API_KEY",
       ai: {
         model: "gemini-1.5-pro",
-        temperature: 0
+        temperature: 0,
       },
       options: {
         debug: true,
-        googleSearchRetrieval: true
+        googleSearchRetrieval: {
+          dynamic_retrieval_config: {
+            mode: "MODE_UNSPECIFIED",
+          },
+        },
       },
-      functions: ["CurrentDateTime", "DaysBetweenDates"]
     },
     {
       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\"",
+      signature:
+        'question:string "a question from a user", plan:string "a suggested plan to answer the question" -> answer:string "the answer"',
       provider: "openai",
       providerKeyName: "OPENAI_API_KEY",
       ai: {
         model: "gpt-4o-mini",
-        temperature: 0
+        temperature: 0,
       },
       options: {
-        debug: true
+        debug: true,
       },
-      functions: ["CurrentDateTime", "DaysBetweenDates"]
-    }
-  ]
+      functions: ["CurrentDateTime", "DaysBetweenDates"],
+    },
+  ],
 };
 
 // Create crew instance with the direct configuration
@@ -81,16 +89,16 @@ const calculateAndPrintAgentCost = (agent, agentName) => {
   console.log(`\n${agentName} Usage:\nPrompt Token Cost: $${(promptTokens / 1000000 * modelDetails?.promptTokenCostPer1M).toFixed(6)}\nCompletion Token Cost: $${(completionTokens / 1000000 * modelDetails?.completionTokenCostPer1M).toFixed(6)}\nTotal Cost: $${totalCost.toFixed(6)}\nPrompt Tokens: ${promptTokens}\nCompletion Tokens: ${completionTokens}\nTotal Tokens: ${promptTokens + completionTokens}`);
 };
 
-const userQuery = "what is the cube root of the number of days between now and Christmas";
+const userQuery = "best sushi in mumbai";
 console.log(`\n\nQuestion: ${userQuery}`);
 
-const Planner = crew.agents.get("Planner");
-const plannerResponse = await Planner.forward({ task: userQuery });
-console.log(`\n\nPlanner Response: ${plannerResponse.plan}`);
+// const Planner = crew.agents.get("Planner");
+// const plannerResponse = await Planner.forward({ task: userQuery });
+// console.log(`\n\nPlanner Response: ${plannerResponse.plan}`);
 
-const Manager = crew.agents.get("Manager");
-const managerResponse = await Manager.forward({ question: userQuery, plan: plannerResponse.plan });
-console.log(`\n\nManager Response: ${managerResponse.answer}`);
+const WebSearch = crew.agents.get("WebSearch");
+const webSearchResponse = await WebSearch.forward({ webSearchQuery: userQuery });
+console.log(`\n\nWeb Search Response: ${webSearchResponse.webSearchResponse}`);
 
 // const calculatorResponse = await crew.agents.get("Calculator").forward({ mathProblem: userQuery });
 // calculateAndPrintAgentCost(crew.agents.get("Calculator"), "Calculator");

package/tsconfig.json

@@ -1,5 +1,6 @@
 {
   "compilerOptions": {
+    "declaration": true,
     "module": "ESNext",
     "target": "ES2020",
     "moduleResolution": "node",