@amitdeshmukh/ax-crew 1.0.0

package/.env.example

@@ -0,0 +1,3 @@
+# API Keys
+OPENAI_API_KEY=sk-******************BlLB
+ANTHROPIC_API_KEY=sk-ant-api********cQAA

package/README.md

@@ -0,0 +1,87 @@
+# Agents
+
+This repo is designed to be used as a module. It simplifies development of AI Agents via config rather than code, so that they can be either reused or rapidly deployed to different applications and processes. The core functionality revolves around the use of AI models to process and respond to queries.
+
+## Getting Started
+
+### Installation
+```bash
+npm install @buddhic-ai/agents
+```
+
+### 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.
+
+## Crew Configuration
+A Crew is a team of agents that work together to achieve a common goal. The configuration file for a crew is a YAML file that defines the agents in the crew, along with their configuration.
+
+An example configuration is provided in `agent_config.example.yaml`. 
+
+### Creating the Crew
+To initialize a crew of agents, pass the path to the configuration file to the `AgentCrew` class.
+
+```javascript
+// Import the AgentCrew class
+import AgentCrew from './src/agents/index.js';
+
+// Create a new instance of AgentCrew
+const configFilePath = './agent_config.example.yaml';
+const crew = new AgentCrew(configFilePath);
+```
+### Adding Agents to the Crew
+You can add a sub-set of defined agents from the configuration file to the crew by passing their names as an array to the `addAgents` method.
+
+Please ensure that the agents are defined in the configuration file before adding them to the crew. Also, the order in which the agents are added to the crew is important, as an error will be thrown if an agent is added before its dependent agents.
+
+For example, the `Manager` agent in the configuration file depends on the `Planner` and `Calculator` agents. Therefore, the `Planner` and `Calculator` agents must be added to the crew before the `Manager` agent.
+
+```javascript
+// Add agents by providing their names
+const agentNames = ['Planner', 'Calculator', 'Manager'];
+const agents = crew.addAgentsToCrew(agentNames);
+
+// Get agent instances
+const Planner = agents.Planner;
+const Manager = agents.Manager;
+```
+
+### State Management
+
+The `StatefulAxAgent` class in `src/agents/index.js` allows for shared state functionality across agents. Sub-agents can be added to an agent to create complex behaviors. All agents in the crew have access to the shared state. State can also be shared with functions that are passed to the agents. To do this, pass the `state` object as an argument to the function class as shown here https://axllm.dev/guides/functions-1/
+
+
+```javascript
+// Set some state (key/value) for the crew
+crew.state.set('name', 'Crew1');
+crew.state.set('location', 'Earth');
+
+// Get the state for the crew
+crew.state.get('name'); // 'Crew1'
+crew.state.getAll(); // { name: 'Crew1', location: 'Earth' }
+``` 
+
+State can also be set/get by individual agents in the crew. This state is shared with all agents and functions if passed in through the functions class constructor.
+
+```javascript
+Planner.state.set('plan', 'some plan'); 
+console.log(Manager.state.getAll()); // { name: 'Crew1', location: 'Earth', plan: 'some plan' }
+```
+
+## Completing a task
+
+An example of how to complete a task using the agents is shown below. The `Planner` agent is used to plan the task, and the `Manager` agent is used to execute the task.
+
+```javascript
+const userQuery = "i referred a friend to active and they were hired and started work on 1st july. But i did not receive my referral bonus. what amount should i have received?";
+
+console.log(`\n\nQuestion: ${userQuery}`);
+
+const planResponse = await Planner.forward({ task: userQuery });
+const managerResponse = await Manager.forward({ question: userQuery, plan: planResponse.plan });
+
+const plan = planResponse.plan;
+const answer = managerResponse.answer;
+
+console.log(`\n\nPlan: ${plan}`);
+console.log(`\n\nAnswer: ${answer}`);
+```

package/agent_config.example.yaml

@@ -0,0 +1,41 @@
+# AxLLM Crew config
+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 based on company policies and operations"
+  provider: anthropic
+  provider_key_name: ANTHROPIC_API_KEY
+  ai:
+    model: claude-3-5-sonnet-20240620
+    temperature: 0
+  options:
+    debug: true
+
+- 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"
+  provider: google-gemini
+  provider_key_name: GEMINI_API_KEY
+  ai:
+    model: gemini-1.5-pro
+    temperature: 0
+  options:
+    debug: true
+    codeExecution: true
+
+- name: Manager
+  description: Answers questions about company policies and operations
+  signature: question:string "a question from a user about company policies and operations", plan:string "a suggested plan to answer the question" -> answer:string "the answer"
+  provider: anthropic
+  provider_key_name: ANTHROPIC_API_KEY
+  ai:
+    model: claude-3-5-sonnet-20240620
+    temperature: 0
+  options:
+    debug: true
+  functions:
+    - CurrentDateTime
+    - DaysBetweenDates
+  agents:
+    - Calculator
+

package/dist/agents/agentConfig.js

@@ -0,0 +1,113 @@
+import fs from 'fs';
+import yaml from 'js-yaml';
+import { PROVIDER_API_KEYS } from '../config/index';
+import { functions as importedFunctions } from '../functions/index';
+import { AxAIAnthropic, AxAIOpenAI, AxAIAzureOpenAI, AxAICohere, AxAIDeepSeek, AxAIGoogleGemini, AxAIGroq, AxAIHuggingFace, AxAIMistral, AxAIOllama, AxAITogether } from '@ax-llm/ax';
+// Define a mapping from provider names to their respective constructors
+const AIConstructors = {
+    'anthropic': AxAIAnthropic,
+    'azure-openai': AxAIAzureOpenAI,
+    'cohere': AxAICohere,
+    'deepseek': AxAIDeepSeek,
+    'google-gemini': AxAIGoogleGemini,
+    'groq': AxAIGroq,
+    'huggingFace': AxAIHuggingFace,
+    'mistral': AxAIMistral,
+    'ollama': AxAIOllama,
+    'openai': AxAIOpenAI,
+    'together': AxAITogether
+};
+const functions = importedFunctions;
+/**
+ * Reads the AI parameters from the YAML configuration file.
+ * @param {string} agentConfigFilePath - The path to the agent_config.yaml file.
+ * @returns {Object} The parsed agent configs from the config.yaml file.
+ * @throws Will throw an error if reading the file fails.
+ */
+const parseAgentConfig = (agentConfigFilePath) => {
+    try {
+        const fileContents = fs.readFileSync(agentConfigFilePath, 'utf8');
+        const parsedConfigs = yaml.load(fileContents);
+        return parsedConfigs;
+    }
+    catch (e) {
+        console.error('Error reading agent config file:', e);
+        throw e;
+    }
+};
+/**
+ * Initializes the AI agent using the specified agent name and configuration file path.
+ * 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 {string} agentConfigFilePath - The file path to the YAML configuration for 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.
+ */
+const getAgentConfigParams = (agentName, agentConfigFilePath, state) => {
+    try {
+        // Retrieve the parameters for the specified AI agent from a config file in yaml format
+        const agentConfig = parseAgentConfig(agentConfigFilePath).crew.find(agent => agent.name === agentName);
+        if (!agentConfig) {
+            throw new Error(`AI agent with name ${agentName} is not configured`);
+        }
+        // Get the constructor for the AI agent's provider
+        const AIConstructor = AIConstructors[agentConfig.provider];
+        if (!AIConstructor) {
+            throw new Error(`AI provider ${agentConfig.provider} is not supported. Did you mean '${agentConfig.provider.toLowerCase()}'?`);
+        }
+        // If an API Key property is present, get the API key for the AI agent from the environment variables
+        let apiKey = '';
+        if (agentConfig.provider_key_name) {
+            apiKey = PROVIDER_API_KEYS[agentConfig.provider_key_name] || '';
+            if (!apiKey) {
+                throw new Error(`API key for provider ${agentConfig.provider} is not set in environment variables`);
+            }
+        }
+        else {
+            throw new Error(`Provider key name is missing in the agent configuration`);
+        }
+        // Create an instance of the AI agent
+        const ai = new AIConstructor({
+            apiKey,
+            config: agentConfig.ai,
+            options: {
+                debug: agentConfig.debug || false
+            }
+        });
+        // If an apiURL is provided in the agent config, set it in the AI agent
+        if (agentConfig.apiURL) {
+            ai.setAPIURL(agentConfig.apiURL);
+        }
+        // Set all options from the agent configuration
+        ai.setOptions({ ...agentConfig.options });
+        // Prepare functions for the AI agent
+        const agentFunctions = (agentConfig.functions || [])
+            .map(funcName => {
+            const func = functions[funcName];
+            if (!func) {
+                console.warn(`Warning: Function ${funcName} not found.`);
+                return;
+            }
+            return func;
+        })
+            .filter(Boolean);
+        // Return AI instance and Agent parameters
+        return {
+            ai,
+            name: agentName,
+            description: agentConfig.description,
+            signature: agentConfig.signature,
+            functions: agentFunctions,
+            subAgentNames: agentConfig.agents || []
+        };
+    }
+    catch (error) {
+        console.error(error);
+        throw new Error(`Error setting up AI agent`);
+    }
+};
+export { getAgentConfigParams };

package/dist/agents/index.js

@@ -0,0 +1,93 @@
+import { v4 as uuidv4 } from 'uuid';
+import { getAgentConfigParams } from './agentConfig';
+import { AxAgent } from '@ax-llm/ax';
+import { createState } from '../state/createState';
+// 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)
+        };
+        super(ai, formattedOptions);
+        this.state = state;
+    }
+}
+/**
+ * Represents a crew of agents with shared state functionality.
+ */
+class AgentCrew {
+    /**
+     * Creates an instance of AgentCrew.
+     * @param {string} configFilePath - The file path to the agent configuration.
+     * @param {string} [crewId=uuidv4()] - The unique identifier for the crew.
+     */
+    constructor(configFilePath, crewId = uuidv4()) {
+        /**
+         * 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.
+         */
+        this.createAgent = (agentName) => {
+            try {
+                const agentConfigParams = getAgentConfigParams(agentName, this.configFilePath, this.state);
+                // Destructure with type assertion
+                const { ai, name, description, signature, functions, subAgentNames } = agentConfigParams;
+                // Get subagents for the AI agent
+                const subAgents = subAgentNames.map((subAgentName) => {
+                    if (!this.agents?.get(subAgentName)) {
+                        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 !== undefined),
+                    agents: subAgents.filter((agent) => agent !== undefined)
+                }, this.state);
+                return agent;
+            }
+            catch (error) {
+                console.error(`Failed to create agent '${agentName}':`, error);
+                throw error;
+            }
+        };
+        this.configFilePath = configFilePath;
+        this.crewId = crewId;
+        this.agents = new Map();
+        this.state = createState(crewId);
+    }
+    /**
+     * Adds an agent to the crew by name.
+     * @param {string} agentName - The name of the agent to add.
+     */
+    addAgent(agentName) {
+        if (this.agents && !this.agents.has(agentName)) {
+            this.agents.set(agentName, this.createAgent(agentName));
+        }
+    }
+    /**
+     * Sets up agents by their names.
+     * For each name provided, it will add the agent to the crew if it's not already present.
+     * @param {string[]} agentNames - An array of agent names to set up.
+     * @returns {Map<string, StatefulAxAgent> | null} A map of agent names to their corresponding instances.
+     */
+    setupAgents(agentNames) {
+        agentNames.forEach(agentName => {
+            this.addAgent(agentName);
+        });
+        return this.agents;
+    }
+    /**
+     * Cleans up the crew by dereferencing agents and resetting the state.
+     */
+    destroy() {
+        this.agents = null;
+        this.state.reset();
+    }
+}
+export default AgentCrew;

package/dist/config/index.js

@@ -0,0 +1,24 @@
+import dotenv from 'dotenv';
+dotenv.config();
+const DEBUG = process.env.DEBUG === 'true';
+// Weaviate
+const WEAVIATE_SCHEME = process.env.WEAVIATE_SCHEME || 'http';
+const WEAVIATE_HOST = process.env.WEAVIATE_HOST || 'localhost:8080';
+const WEAVIATE_APIKEY = process.env.WEAVIATE_APIKEY || '';
+// AI API keys
+const ANTHROPIC_API_KEY = process.env.ANTHROPIC_API_KEY;
+const OPENAI_API_KEY = process.env.OPENAI_API_KEY;
+const COHERE_API_KEY = process.env.COHERE_API_KEY;
+const GEMINI_API_KEY = process.env.GEMINI_API_KEY;
+const PROVIDER_API_KEYS = {
+    COHERE_API_KEY,
+    GEMINI_API_KEY,
+    OPENAI_API_KEY,
+    ANTHROPIC_API_KEY,
+};
+// GCP
+const GCP_PROJECT_ID = process.env.GCP_PROJECT_ID;
+const GCP_LOCATION = process.env.GCP_LOCATION;
+const GCP_RAG_KEYRING = process.env.GCP_RAG_KEYRING;
+const GCP_CRYPTO_SECRET = process.env.GCP_CRYPTO_SECRET;
+export { DEBUG, WEAVIATE_SCHEME, WEAVIATE_HOST, WEAVIATE_APIKEY, PROVIDER_API_KEYS, GCP_PROJECT_ID, GCP_LOCATION, GCP_RAG_KEYRING, GCP_CRYPTO_SECRET };

package/dist/functions/dateTime.js

@@ -0,0 +1,36 @@
+export const CurrentDateTime = {
+    name: 'CurrentDateTime',
+    description: 'Get the current date and time.',
+    parameters: {
+        type: 'object',
+        properties: {},
+        required: []
+    },
+    func: () => new Date().toISOString()
+};
+export const DaysBetweenDates = {
+    name: 'DaysBetweenDates',
+    description: 'Calculate the number of days between two dates in ISO format.',
+    parameters: {
+        type: 'object',
+        properties: {
+            startDate: {
+                type: 'string',
+                description: 'The start date in ISO format.'
+            },
+            endDate: {
+                type: 'string',
+                description: 'The end date in ISO format.'
+            }
+        },
+        required: ['startDate', 'endDate']
+    },
+    func: (args) => {
+        const { startDate, endDate } = args;
+        const start = new Date(startDate);
+        const end = new Date(endDate);
+        const differenceInTime = end.getTime() - start.getTime();
+        const differenceInDays = differenceInTime / (1000 * 3600 * 24);
+        return Math.round(differenceInDays);
+    }
+};

package/dist/functions/index.js

@@ -0,0 +1,6 @@
+import { CurrentDateTime, DaysBetweenDates } from './dateTime';
+const functions = {
+    CurrentDateTime,
+    DaysBetweenDates
+};
+export { functions };

package/dist/index.js

@@ -0,0 +1,2 @@
+import AgentCrew from './agents/index';
+export { AgentCrew };

package/dist/state/createState.js

@@ -0,0 +1,32 @@
+const stateInstances = {};
+const createState = (id) => {
+    if (!id) {
+        throw new Error('An ID is required to create a new state instance.');
+    }
+    if (stateInstances[id]) {
+        return stateInstances[id];
+    }
+    let data = {};
+    stateInstances[id] = {
+        set: (key, value) => {
+            data[key] = value;
+        },
+        get: (key) => {
+            return data[key];
+        },
+        getAll: () => {
+            return data;
+        },
+        getId: () => {
+            return id;
+        },
+        reset: () => {
+            data = {};
+        }
+    };
+    return stateInstances[id];
+};
+const getState = (id) => {
+    return stateInstances[id];
+};
+export { createState, getState };

package/dist/utils/index.js

@@ -0,0 +1,25 @@
+/**
+ * Checks if a given function is a class constructor.
+ *
+ * @param {Function} obj - The object to check.
+ * @returns {boolean} True if the object is a class constructor, false otherwise.
+ */
+export function isClass(obj) {
+    // A class must be a function
+    if (typeof obj !== 'function') {
+        return false;
+    }
+    // The function must have a prototype, but not be an instance of another function
+    if (obj.prototype === undefined || obj.prototype.constructor !== obj) {
+        return false;
+    }
+    // The function's prototype must have isPrototypeOf method
+    if (!Object.prototype.hasOwnProperty.call(obj.prototype, 'isPrototypeOf')) {
+        return false;
+    }
+    // The function must not have its own 'isPrototypeOf' property
+    if (obj.hasOwnProperty('isPrototypeOf')) {
+        return false;
+    }
+    return true;
+}

package/package.json

@@ -0,0 +1,27 @@
+{
+  "type": "module",
+  "name": "@amitdeshmukh/ax-crew",
+  "version": "1.0.0",
+  "description": "Build and launch a crew of AI agents with shared state. Built with axllm.dev",
+  "main": "dist/index.js",
+  "engines": {
+    "node": ">=21.0.0"
+  },
+  "scripts": {
+    "build": "tsc --outDir dist",
+    "release": "npm run build && npm publish --access public"
+  },
+  "dependencies": {
+    "@ax-llm/ax": "^9.0.23",
+    "dotenv": "^16.4.5",
+    "js-yaml": "^4.1.0",
+    "uuid": "^10.0.0"
+  },
+  "license": "MIT",
+  "devDependencies": {
+    "@types/js-yaml": "^4.0.9",
+    "@types/node": "^20.14.9",
+    "@types/uuid": "^10.0.0",
+    "typescript": "^5.5.3"
+  }
+}

package/src/agents/agentConfig.ts

@@ -0,0 +1,157 @@
+import fs from 'fs';
+import yaml from 'js-yaml';
+import { PROVIDER_API_KEYS } from '../config/index';
+import { functions as importedFunctions } from '../functions/index';
+
+import { AxAIAnthropic, AxAIOpenAI, AxAIAzureOpenAI, AxAICohere, AxAIDeepSeek, AxAIGoogleGemini, AxAIGroq, AxAIHuggingFace, AxAIMistral, AxAIOllama, AxAITogether } from '@ax-llm/ax';
+import type { AxAI, AxModelConfig, AxFunction, AxSignature } from '@ax-llm/ax';
+
+// Define a mapping from provider names to their respective constructors
+const AIConstructors: Record<string, any> = {
+  'anthropic': AxAIAnthropic,
+  'azure-openai': AxAIAzureOpenAI,
+  'cohere': AxAICohere,
+  'deepseek': AxAIDeepSeek,
+  'google-gemini': AxAIGoogleGemini,
+  'groq': AxAIGroq,
+  'huggingFace': AxAIHuggingFace,
+  'mistral': AxAIMistral,
+  'ollama': AxAIOllama,
+  'openai': AxAIOpenAI,
+  'together': AxAITogether
+};
+
+type ExtendedAxModelConfig = AxModelConfig & {
+  model: string;
+};
+
+// Define a mapping from function names to their respective handlers
+type FunctionMap = {
+  [key: string]: AxFunction | { toFunction: () => AxFunction };
+};
+const functions: FunctionMap = importedFunctions;
+
+interface AgentConfig {
+  name: string;
+  description: string;
+  signature: AxSignature;
+  provider: string;
+  provider_key_name?: string;
+  ai: ExtendedAxModelConfig;
+  debug?: boolean;
+  apiURL?: string;
+  options?: Record<string, any>;
+  functions?: string[];
+  agents?: string[];
+}
+
+interface Agent {
+  ai: AxAI;
+  name: string;
+  description: string;
+  signature?: AxSignature;
+  functions: AxFunction[];
+  subAgentNames: string[];
+}
+
+
+/**
+ * Reads the AI parameters from the YAML configuration file.
+ * @param {string} agentConfigFilePath - The path to the agent_config.yaml file.
+ * @returns {Object} The parsed agent configs from the config.yaml file.
+ * @throws Will throw an error if reading the file fails.
+ */
+const parseAgentConfig = (agentConfigFilePath: string): {crew: AgentConfig[]} => {
+  try {
+    const fileContents = fs.readFileSync(agentConfigFilePath, 'utf8');
+    const parsedConfigs = yaml.load(fileContents) as { crew: AgentConfig[] };
+    return parsedConfigs;
+  } catch (e) {
+    console.error('Error reading agent config file:', e);
+    throw e;
+  }
+};
+
+/**
+ * Initializes the AI agent using the specified agent name and configuration file path.
+ * 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 {string} agentConfigFilePath - The file path to the YAML configuration for 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.
+ */
+const getAgentConfigParams = (agentName: string, agentConfigFilePath: string, state: Record<string, any>) => {
+  try{
+    // Retrieve the parameters for the specified AI agent from a config file in yaml format
+    const agentConfig = parseAgentConfig(agentConfigFilePath).crew.find(agent => agent.name === agentName);
+    if (!agentConfig) {
+      throw new Error(`AI agent with name ${agentName} is not configured`);
+    }
+
+    // Get the constructor for the AI agent's provider
+    const AIConstructor = AIConstructors[agentConfig.provider];
+    if (!AIConstructor) {
+      throw new Error(`AI provider ${agentConfig.provider} is not supported. Did you mean '${agentConfig.provider.toLowerCase()}'?`);
+    }
+
+    // If an API Key property is present, get the API key for the AI agent from the environment variables
+    let apiKey = '';
+    if (agentConfig.provider_key_name) {
+      apiKey = PROVIDER_API_KEYS[agentConfig.provider_key_name] || '';
+
+      if (!apiKey) {
+        throw new Error(`API key for provider ${agentConfig.provider} is not set in environment variables`);
+      }
+    } else {
+      throw new Error(`Provider key name is missing in the agent configuration`);
+    }
+
+    // Create an instance of the AI agent
+    const ai = new AIConstructor({
+      apiKey,
+      config: agentConfig.ai,
+      options: {
+        debug: agentConfig.debug || false
+      }
+    });
+
+    // If an apiURL is provided in the agent config, set it in the AI agent
+    if (agentConfig.apiURL) {
+      ai.setAPIURL(agentConfig.apiURL);
+    }
+
+    // Set all options from the agent configuration
+    ai.setOptions({ ...agentConfig.options });
+    
+    // Prepare functions for the AI agent
+    const agentFunctions = (agentConfig.functions || [])
+      .map(funcName => {
+        const func = functions[funcName];
+        if (!func) {
+          console.warn(`Warning: Function ${funcName} not found.`);
+          return;
+        }        
+        return func;
+      })
+      .filter(Boolean);
+
+    // Return AI instance and Agent parameters
+    return {
+      ai,
+      name: agentName,
+      description: agentConfig.description,
+      signature: agentConfig.signature,
+      functions: agentFunctions,
+      subAgentNames: agentConfig.agents || []
+    };
+  } catch (error) {
+    console.error(error);
+    throw new Error(`Error setting up AI agent`);
+  }
+};
+
+export { getAgentConfigParams }

package/src/agents/index.ts

@@ -0,0 +1,133 @@
+import { v4 as uuidv4 } from 'uuid';
+import { getAgentConfigParams } from './agentConfig';
+import { AxAgent, AxAI } from '@ax-llm/ax';
+import type { AxSignature, AxAgentic, AxFunction, AxAIService } from '@ax-llm/ax';
+import { createState, StateInstance } from '../state/createState';
+
+
+// Define interfaces for the agent configurations and return types
+interface Agents {
+  [key: string]: StatefulAxAgent;
+}
+
+interface AgentConfigParams {
+  ai: AxAI;
+  name: string;
+  description: string;
+  signature: AxSignature;
+  functions: (AxFunction | { toFunction: () => AxFunction; } | undefined)[];
+  subAgentNames: string[];
+}
+
+// Extend the AxAgent class to include shared state functionality
+class StatefulAxAgent extends AxAgent<any, any> {
+  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
+    };
+    super(ai, formattedOptions);
+    this.state = state;
+  }
+}
+
+
+/**
+ * Represents a crew of agents with shared state functionality.
+ */
+class AgentCrew {
+  private configFilePath: string;
+  crewId: string;
+  agents: Map<string, StatefulAxAgent> | null;
+  state: StateInstance;
+
+  /**
+   * Creates an instance of AgentCrew.
+   * @param {string} configFilePath - The file path to the agent configuration.
+   * @param {string} [crewId=uuidv4()] - The unique identifier for the crew.
+   */
+  constructor(configFilePath: string, crewId: string = uuidv4()) {
+    this.configFilePath = configFilePath;
+    this.crewId = crewId;
+    this.agents = new Map<string, StatefulAxAgent>();
+    this.state = createState(crewId);
+  }
+
+  /**
+   * 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 => {
+    try {
+      const agentConfigParams: AgentConfigParams = getAgentConfigParams(agentName, this.configFilePath, this.state);
+
+      // Destructure with type assertion
+      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.`);
+        }
+        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);
+      
+      return agent;
+    } catch (error) {
+      console.error(`Failed to create agent '${agentName}':`, error);
+      throw error;
+    }
+  };
+
+  /**
+   * Adds an agent to the crew by name.
+   * @param {string} agentName - The name of the agent to add.
+   */
+  addAgent(agentName: string): void {
+    if (this.agents && !this.agents.has(agentName)) {
+      this.agents.set(agentName, this.createAgent(agentName));
+    }
+  }  
+
+  /**
+   * Sets up agents by their names.
+   * For each name provided, it will add the agent to the crew if it's not already present.
+   * @param {string[]} agentNames - An array of agent names to set up.
+   * @returns {Map<string, StatefulAxAgent> | null} A map of agent names to their corresponding instances.
+   */
+  setupAgents(agentNames: string[]): Map<string, StatefulAxAgent> | null {
+    agentNames.forEach(agentName => {
+      this.addAgent(agentName);
+    });
+    return this.agents;
+  }
+
+  /**
+   * Cleans up the crew by dereferencing agents and resetting the state.
+   */
+  destroy() {
+    this.agents = null;
+    this.state.reset();
+  }  
+}
+
+export default AgentCrew;

package/src/config/index.ts

@@ -0,0 +1,44 @@
+import dotenv from 'dotenv';
+dotenv.config();
+
+const DEBUG: boolean = process.env.DEBUG === 'true';
+
+// Weaviate
+const WEAVIATE_SCHEME: string = process.env.WEAVIATE_SCHEME || 'http';
+const WEAVIATE_HOST: string = process.env.WEAVIATE_HOST || 'localhost:8080';
+const WEAVIATE_APIKEY: string = process.env.WEAVIATE_APIKEY || '';
+
+// AI API keys
+const ANTHROPIC_API_KEY: string | undefined = process.env.ANTHROPIC_API_KEY;
+const OPENAI_API_KEY: string | undefined = process.env.OPENAI_API_KEY;
+const COHERE_API_KEY: string | undefined = process.env.COHERE_API_KEY;
+const GEMINI_API_KEY: string | undefined = process.env.GEMINI_API_KEY;
+
+interface ProviderApiKeys {
+  [key: string]: string | undefined;
+}
+
+const PROVIDER_API_KEYS: ProviderApiKeys = {
+  COHERE_API_KEY,
+  GEMINI_API_KEY,
+  OPENAI_API_KEY,
+  ANTHROPIC_API_KEY,
+};
+
+// GCP
+const GCP_PROJECT_ID: string | undefined = process.env.GCP_PROJECT_ID;
+const GCP_LOCATION: string | undefined = process.env.GCP_LOCATION;
+const GCP_RAG_KEYRING: string | undefined = process.env.GCP_RAG_KEYRING;
+const GCP_CRYPTO_SECRET: string | undefined = process.env.GCP_CRYPTO_SECRET;
+
+export {
+  DEBUG,
+  WEAVIATE_SCHEME,
+  WEAVIATE_HOST,
+  WEAVIATE_APIKEY,
+  PROVIDER_API_KEYS,
+  GCP_PROJECT_ID,
+  GCP_LOCATION,
+  GCP_RAG_KEYRING,
+  GCP_CRYPTO_SECRET
+};

package/src/functions/dateTime.ts

@@ -0,0 +1,39 @@
+import type { AxFunction } from "@ax-llm/ax";
+
+export const CurrentDateTime: AxFunction = {
+  name: 'CurrentDateTime',
+  description: 'Get the current date and time.',
+  parameters: {
+    type: 'object',
+    properties: {},
+    required: []
+  },
+  func: () => new Date().toISOString()
+}
+
+export const DaysBetweenDates: AxFunction = {
+  name: 'DaysBetweenDates',
+  description: 'Calculate the number of days between two dates in ISO format.',
+  parameters: {
+    type: 'object',
+    properties: {
+      startDate: { 
+        type: 'string',
+        description: 'The start date in ISO format.'
+      },
+      endDate: {
+        type: 'string',
+        description: 'The end date in ISO format.'
+      }
+    },
+    required: ['startDate', 'endDate']
+  },
+  func: (args: { startDate: string; endDate: string }) => {
+    const { startDate, endDate } = args;
+    const start = new Date(startDate);
+    const end = new Date(endDate);
+    const differenceInTime = end.getTime() - start.getTime();
+    const differenceInDays = differenceInTime / (1000 * 3600 * 24);
+    return Math.round(differenceInDays);
+  }
+}

package/src/functions/index.ts

@@ -0,0 +1,8 @@
+import { CurrentDateTime, DaysBetweenDates } from './dateTime';
+
+const functions = {
+  CurrentDateTime,
+  DaysBetweenDates
+}
+
+export { functions };

package/src/index.ts

@@ -0,0 +1,3 @@
+import AgentCrew from './agents/index';
+
+export { AgentCrew };

package/src/state/createState.ts

@@ -0,0 +1,46 @@
+interface StateInstance {
+  set: (key: string, value: any) => void;
+  get: (key: string) => any;
+  getAll: () => { [key: string]: any };
+  getId: () => string;
+  reset: () => void;
+}
+
+const stateInstances: { [id: string]: StateInstance } = {};
+
+const createState = (id: string): StateInstance => {
+  if (!id) {
+    throw new Error('An ID is required to create a new state instance.');
+  }
+  if (stateInstances[id]) {
+    return stateInstances[id];
+  }
+
+  let data: { [key: string]: any } = {};
+
+  stateInstances[id] = {
+    set: (key, value) => {
+      data[key] = value;
+    },
+    get: (key) => {
+      return data[key];
+    },
+    getAll: () => {
+      return data;
+    },
+    getId: () => {
+      return id;
+    },
+    reset: () => {
+      data = {};
+    }
+  };
+
+  return stateInstances[id];
+};
+
+const getState = (id: string): StateInstance | undefined  => {
+  return stateInstances[id];
+}
+
+export { createState, getState, StateInstance };

package/src/utils/index.ts

@@ -0,0 +1,27 @@
+import { AxFunction } from "@ax-llm/ax";
+
+/**
+ * Checks if a given function is a class constructor.
+ *
+ * @param {Function} obj - The object to check.
+ * @returns {boolean} True if the object is a class constructor, false otherwise.
+ */
+export function isClass(obj: AxFunction | ClassDecorator): boolean {
+  // A class must be a function
+  if (typeof obj !== 'function') {
+    return false;
+  }
+  // The function must have a prototype, but not be an instance of another function
+  if (obj.prototype === undefined || obj.prototype.constructor !== obj) {
+    return false;
+  }
+  // The function's prototype must have isPrototypeOf method
+  if (!Object.prototype.hasOwnProperty.call(obj.prototype, 'isPrototypeOf')) {
+    return false;
+  }
+  // The function must not have its own 'isPrototypeOf' property
+  if (obj.hasOwnProperty('isPrototypeOf')) {
+    return false;
+  }
+  return true;
+}

package/test.js

@@ -0,0 +1,36 @@
+// Import the AgentCrew class
+import AgentCrew from './dist/agents/index.js';
+
+// Create a new instance of AgentCrew
+const configFilePath = './agent_config.example.yaml';
+const crew = new AgentCrew(configFilePath);
+
+// Add agents by providing their names
+const agentNames = ['Planner', 'Calculator', 'Manager'];
+const agents = crew.addAgentsToCrew(agentNames);
+
+crew.state.set('clientId', '0906-978189-7350');
+console.log(`Set ClientId: ${crew.state.get('clientId')}`);
+
+const ManagerAgent = agents.Manager;
+const PlannerAgent = agents['Planner'];
+
+console.log(crew.state.getAll())
+PlannerAgent.state.set('plan', 'plan 1 2 3')
+console.log(ManagerAgent.state.getAll())
+
+const userQuery = "i referred a friend to active and they were hired and started work on 1st july. But i did not receive my referral bonus. what amount should i have received?";
+
+console.log(`\n\nQuestion: ${userQuery}`);
+
+const planResponse = await PlannerAgent.forward({ task: userQuery });
+const managerResponse = await ManagerAgent.forward({ question: userQuery, plan: planResponse.plan });
+
+const plan = planResponse.plan;
+const answer = managerResponse.answer;
+const reason = managerResponse.reason;
+
+console.log(`\n\nPlan: ${plan}`);
+console.log(`\n\nAnswer: ${answer}`);
+
+console.log(crew.state.getAll());

package/tsconfig.json

@@ -0,0 +1,17 @@
+{
+  "compilerOptions": {
+    "module": "ESNext",
+    "target": "ES2020",
+    "moduleResolution": "node",
+    "esModuleInterop": true,
+    "allowSyntheticDefaultImports": true,
+    "strict": true,
+    "outDir": "./dist",
+    "rootDir": "./src",
+    "allowJs": true,
+    "resolveJsonModule": true,
+    "typeRoots": ["./node_modules/@types", "./types"]
+  },
+  "include": ["src/**/*"],
+  "exclude": ["node_modules", "**/*.spec.ts"]
+}