@amitdeshmukh/ax-crew 3.3.2 → 3.3.4

package/CHANGELOG.md

@@ -5,6 +5,24 @@ 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
+
+### Removed
+- Removed unused `js-yaml` dependency
+- Updated `README.md` to highlight AxLLM 10.0.9 as peer dependency
+
+## [3.3.3] - 2024-12-12
+
+### Added
+- TypeScript declaration file (`index.d.ts`)
+- Improved type definitions for:
+  - `AxCrew` class and its methods
+  - `StatefulAxAgent` class
+  - State management interfaces
+  - Function registry types
+  - Configuration types
+- TypeScript documentation in `README.md` with example
+
 ## [3.3.1] - 2024-12-11
 
 ### Changed

package/README.md

@@ -16,12 +16,94 @@ Install this package:
 ```bash
 npm install @amitdeshmukh/ax-crew
 ```
-AxLLM is a peer dependency, so you will need to install it separately.
+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.
 
 ```bash
-npm install @ax-llm/ax@latest
+npm install @ax-llm/ax@10.0.9
 ```
 
+### TypeScript Support
+This package includes TypeScript declarations and provides full type safety. Here's how to use it with TypeScript:
+
+```typescript
+import { AxCrew, AxCrewFunctions, FunctionRegistryType, StateInstance } from '@amitdeshmukh/ax-crew';
+import type { AxFunction } from '@ax-llm/ax';
+
+// Type-safe configuration
+const 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\"",
+    provider: "google-gemini",
+    providerKeyName: "GEMINI_API_KEY",
+    ai: {
+      model: "gemini-1.5-pro",
+      temperature: 0
+    }
+  }]
+};
+
+// Create custom functions with type safety
+class MyCustomFunction {
+  constructor(private state: Record<string, any>) {}
+  
+  toFunction(): AxFunction {
+    return {
+      name: 'MyCustomFunction',
+      description: 'Does something useful',
+      parameters: {
+        type: 'object',
+        properties: {
+          inputParam: { type: 'string', description: "input to the function" }
+        }
+      },
+      func: async ({ inputParam }) => {
+        // Implementation
+        return inputParam;
+      }
+    };
+  }
+}
+
+// Type-safe function registry
+const myFunctions: FunctionRegistryType = {
+  MyCustomFunction
+};
+
+// Create crew with type checking
+const crew = new AxCrew(config, myFunctions);
+
+// Type-safe state management
+crew.state.set('key', 'value');
+const value: string = crew.state.get('key');
+
+// Type-safe agent management
+const agents = crew.addAgentsToCrew(['Planner']);
+const planner = agents?.get('Planner');
+
+if (planner) {
+  // Type-safe agent usage
+  const response = await planner.forward({ task: "Plan something" });
+  const cost = planner.getUsageCost();
+  
+  if (cost) {
+    console.log(`Total cost: $${cost.totalCost}`);
+    console.log(`Total tokens: ${cost.tokenMetrics.totalTokens}`);
+  }
+}
+```
+
+Key TypeScript features:
+- Full type definitions for all classes, methods, and properties
+- Type-safe configuration objects
+- Proper typing for function registries and custom functions
+- Type checking for state management
+- Comprehensive type safety for agent operations and responses
+- Usage cost tracking with proper types
+
 ### 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.
 

package/axllm.js

@@ -10,7 +10,7 @@ const googleAI = new AxAI({
     temperature: 0,
   },
   options: {
-    debug: true,
+    debug: false,
     codeExecution: true,
   },
 });
@@ -19,20 +19,22 @@ const openAI = new AxAI({
   name: "openai",
   apiKey: process.env.OPENAI_API_KEY,
   config: {
-    model: "gpt-4-turbo",
+    model: "gpt-4o-mini",
     maxTokens: 1000,
     temperature: 0,
   },
-  options: { debug: true },
+  options: { debug: false },
 });
 
 const MathAgent = new AxAgent({
+  ai: googleAI,
   name: "MathAgent",
   description: "Solves math problems",
-  signature: `mathProblem:string "a sentence describing a math problem to be solved using Python code"  -> solution:string "a precise answer to the math problem"`,
+  signature: `mathProblem:string "a sentence describing a math problem to be solved using Python code"  -> solution:string "the answer to the math problem"`
 });
 
 const ManagerAgent = new AxAgent({
+  ai: openAI,
   name: "ManagerAgent",
   description: "Completes a user specified task.",
   signature: `question:string "a question to be answered" -> answer:string "the answer to the question"`,
@@ -50,7 +52,7 @@ const managerAgentResponse = await ManagerAgent.forward(
   }
 );
 
-console.log(`\nManagerAgentResponse: ${JSON.stringify(managerAgentResponse, null, 2)}\n`);
+console.log(`\nAnswer: ${JSON.stringify(managerAgentResponse.answer, null, 2)}\n`);
 
 console.log("❌ This is the wrong answer ☝️");
 
@@ -61,6 +63,6 @@ const MathAgentResponse = await MathAgent.forward(
   }
 );
 
-console.log(`\nMathAgentResponse: ${JSON.stringify(MathAgentResponse, null, 2)}\n`);
+console.log(`\nAnswer: ${JSON.stringify(MathAgentResponse.solution, null, 2)}\n`);
 console.log(" ✅ Correct answer ☝️");
 

package/dist/agents/agentConfig.d.ts

@@ -1,15 +1,14 @@
 import type { AxModelConfig, AxFunction, AxSignature } from '@ax-llm/ax';
 import { FunctionRegistryType } from '../functions/index.js';
-type ExtendedAxModelConfig = AxModelConfig & {
-    model: string;
-};
-interface AgentConfig {
+export interface AgentConfig {
     name: string;
     description: string;
-    signature: AxSignature;
+    signature: string | AxSignature;
     provider: string;
     providerKeyName?: string;
-    ai: ExtendedAxModelConfig;
+    ai: AxModelConfig & {
+        model: string;
+    };
     debug?: boolean;
     apiURL?: string;
     options?: Record<string, any>;
@@ -37,7 +36,7 @@ declare const getAgentConfigParams: (agentName: string, agentConfig: AgentConfig
     ai: any;
     name: string;
     description: string;
-    signature: AxSignature;
+    signature: string | AxSignature;
     functions: (AxFunction | (new (state: Record<string, any>) => {
         toFunction: () => AxFunction;
     }) | undefined)[];

package/dist/agents/index.js

@@ -16,7 +16,7 @@ class StatefulAxAgent extends AxAgent {
         this.axai = ai;
         // Set examples if provided
         if (examples && examples.length > 0) {
-            this.setExamples(examples);
+            super.setExamples(examples);
         }
     }
     // Override the forward method

package/dist/index.d.ts

@@ -1,3 +1,6 @@
 import { AxCrew } from './agents/index.js';
 import { AxCrewFunctions, FunctionRegistryType } from './functions/index.js';
-export { AxCrew, AxCrewFunctions, FunctionRegistryType };
+import type { UsageCost } from './agents/agentUseCosts.js';
+import type { AgentConfig, AgentConfigInput } from './agents/agentConfig.js';
+import type { StateInstance } from './state/index.js';
+export { AxCrew, AxCrewFunctions, FunctionRegistryType, type UsageCost, type AgentConfig, type AgentConfigInput, type StateInstance };

package/dist/state/index.d.ts

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

package/dist/state/index.js

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

package/index.d.ts

@@ -1,7 +1,92 @@
+import { AxAI, AxAgentic, AxFunction, AxSignature, AxProgramForwardOptions, AxModelConfig, AxAgent } from "@ax-llm/ax";
+
+export interface UsageCost {
+  promptCost: number;
+  completionCost: number;
+  totalCost: number;
+  tokenMetrics: {
+    promptTokens: number;
+    completionTokens: number;
+    totalTokens: number;
+  };
+}
+
+export interface StateInstance {
+  reset(): void;
+  set(key: string, value: any): void;
+  get(key: string): any;
+  getAll(): Record<string, any>;
+  [key: string]: any;
+}
+
+export type FunctionRegistryType = {
+  [key: string]: AxFunction | { new(state: Record<string, any>): { toFunction: () => AxFunction } };
+};
+
+export interface AgentConfig {
+  name: string;
+  description: string;
+  signature: AxSignature;
+  provider: string;
+  providerKeyName?: string;
+  ai: AxModelConfig & { model: string };
+  debug?: boolean;
+  apiURL?: string;
+  options?: Record<string, any>;
+  functions?: string[];
+  agents?: string[];
+  examples?: Array<Record<string, any>>;
+}
+
+export type AgentConfigInput = string | { crew: AgentConfig[] };
+
+export class StatefulAxAgent extends AxAgent<Record<string, any>, Record<string, any>> {
+  state: StateInstance;
+  axai: AxAI;
+
+  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;
+
+  setExamples(examples: Array<Record<string, any>>): void;
+}
+
 export class AxCrew {
-  constructor(configPath: string);
-  addAgent(agent: any): void;
-  createAgent(agent: any): void;
-  addAgentsToCrew(agents: any[]): void;
+  private agentConfig: AgentConfigInput;
+  functionsRegistry: FunctionRegistryType;
+  crewId: string;
+  agents: Map<string, StatefulAxAgent> | null;
+  state: StateInstance;
+
+  constructor(
+    agentConfig: AgentConfigInput,
+    functionsRegistry?: FunctionRegistryType,
+    crewId?: string
+  );
+
+  createAgent(agentName: string): StatefulAxAgent;
+  addAgent(agentName: string): void;
+  addAgentsToCrew(agentNames: string[]): Map<string, StatefulAxAgent> | null;
   destroy(): void;
 }
+
+export const AxCrewFunctions: {
+  CurrentDateTime: { new(state: Record<string, any>): { toFunction: () => AxFunction } };
+  DaysBetweenDates: { new(state: Record<string, any>): { toFunction: () => AxFunction } };
+};

package/package.json

@@ -1,7 +1,7 @@
 {
   "type": "module",
   "name": "@amitdeshmukh/ax-crew",
-  "version": "3.3.2",
+  "version": "3.3.4",
   "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",
@@ -14,15 +14,13 @@
   },
   "dependencies": {
     "dotenv": "^16.4.5",
-    "js-yaml": "^4.1.0",
     "upgrade": "^1.1.0",
     "uuid": "^10.0.0"
   },
   "peerDependencies": {
-    "@ax-llm/ax": "^10.0.9"
+    "@ax-llm/ax": "10.0.9"
   },
   "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

@@ -20,17 +20,15 @@ const AIConstructors: Record<string, any> = {
   'together': AxAITogether
 };
 
-type ExtendedAxModelConfig = AxModelConfig & {
-  model: string;
-};
 
-interface AgentConfig {
+
+export interface AgentConfig {
   name: string;
   description: string;
-  signature: AxSignature;
+  signature: string | AxSignature;
   provider: string;
   providerKeyName?: string;
-  ai: ExtendedAxModelConfig;
+  ai: AxModelConfig & { model: string };
   debug?: boolean;
   apiURL?: string;
   options?: Record<string, any>;

package/src/agents/index.ts

@@ -17,7 +17,7 @@ interface AgentConfigParams {
   ai: AxAI;
   name: string;
   description: string;
-  signature: AxSignature;
+  signature: string | AxSignature;
   functions: (
     | AxFunction
     | (new (state: Record<string, any>) => { toFunction: () => AxFunction })
@@ -57,7 +57,7 @@ class StatefulAxAgent extends AxAgent<any, any> {
 
     // Set examples if provided
     if (examples && examples.length > 0) {
-      this.setExamples(examples);
+      super.setExamples(examples);
     }
   }
 

package/src/index.ts

@@ -1,8 +1,16 @@
 import { AxCrew } from './agents/index.js';
 import { AxCrewFunctions, FunctionRegistryType } from './functions/index.js';
+import type { UsageCost } from './agents/agentUseCosts.js';
+import type { AgentConfig, AgentConfigInput } from './agents/agentConfig.js';
+import type { StateInstance } from './state/index.js';
 
 export { 
   AxCrew,
   AxCrewFunctions,
-  FunctionRegistryType
+  FunctionRegistryType,
+  // Type exports
+  type UsageCost,
+  type AgentConfig,
+  type AgentConfigInput,
+  type StateInstance
 };

package/src/state/index.ts

@@ -1,46 +1,41 @@
-interface StateInstance {
-  set: (key: string, value: any) => void;
-  get: (key: string) => any;
-  getAll: () => { [key: string]: any };
-  getId: () => string;
-  reset: () => void;
+export interface StateInstance {
+  reset(): void;
+  set(key: string, value: any): void;
+  get(key: string): any;
+  getAll(): Record<string, any>;
+  [key: string]: any;
 }
 
 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.');
-  }
+function createState(id: string): StateInstance {
   if (stateInstances[id]) {
     return stateInstances[id];
   }
 
-  let data: { [key: string]: any } = {};
-
-  stateInstances[id] = {
-    set: (key, value) => {
-      data[key] = value;
-    },
-    get: (key) => {
-      return data[key];
+  let state: Record<string, any> = {};
+  
+  const instance: StateInstance = {
+    reset() {
+      state = {};
     },
-    getAll: () => {
-      return data;
+    set(key: string, value: any) {
+      state[key] = value;
     },
-    getId: () => {
-      return id;
+    get(key: string) {
+      return state[key];
     },
-    reset: () => {
-      data = {};
+    getAll() {
+      return { ...state };
     }
   };
 
-  return stateInstances[id];
-};
+  stateInstances[id] = instance;
+  return instance;
+}
 
-const getState = (id: string): StateInstance | undefined  => {
+function getState(id: string): StateInstance | undefined {
   return stateInstances[id];
 }
 
-export { createState, getState, StateInstance };
+export { createState, getState };

package/test1.js

@@ -3,15 +3,29 @@ import { AxCrew, AxCrewFunctions } from "./dist/index.js";
 // Define the configuration object directly
 const 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 in 5 steps or less"',
+    //   provider: "google-gemini",
+    //   providerKeyName: "GEMINI_API_KEY",
+    //   ai: {
+    //     model: "gemini-1.5-flash",
+    //     temperature: 0,
+    //   },
+    //   options: {
+    //     debug: false,
+    //   },
+    // },
     {
-      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: "google-gemini",
-      providerKeyName: "GEMINI_API_KEY",
+      name: "Summarizer",
+      description: "Summarizes long form text into a short summary",
+      signature: "inputText:string -> summary:string",
+      provider: "openai",
+      providerKeyName: "OPENAI_API_KEY",
       ai: {
-        model: "gemini-1.5-flash",
+        model: "gpt-4o-mini",
         temperature: 0,
       },
       options: {
@@ -19,28 +33,9 @@ const config = {
       },
     },
     {
-      name: "Calculator",
+      name: "MathAgent",
       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",
-      providerKeyName: "GEMINI_API_KEY",
-      ai: {
-        model: "gemini-1.5-pro",
-        temperature: 0,
-      },
-      options: {
-        debug: true,
-        codeExecution: true,
-      },
-      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"',
+      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: {
@@ -48,19 +43,35 @@ const config = {
         temperature: 0,
       },
       options: {
-        debug: true,
-        googleSearchRetrieval: {
-          dynamic_retrieval_config: {
-            mode: "MODE_UNSPECIFIED",
-          },
-        },
+        debug: false,
+        stream: false,
+        codeExecution: true
       },
     },
+    // {
+    //   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"',
+    //   provider: "google-gemini",
+    //   providerKeyName: "GEMINI_API_KEY",
+    //   ai: {
+    //     model: "gemini-1.5-pro",
+    //     temperature: 0,
+    //   },
+    //   options: {
+    //     debug: false,
+    //     stream: false,
+    //     googleSearchRetrieval: {
+    //       mode: "MODE_UNSPECIFIED",
+    //     },
+    //   },
+    // },
     {
-      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"',
+      name: "ManagerAgent",
+      description: "Completes a user specified task.",
+      signature: `task:string "a task to be completed" -> result:string "the result of the task"`,
       provider: "openai",
       providerKeyName: "OPENAI_API_KEY",
       ai: {
@@ -69,16 +80,13 @@ const config = {
       },
       options: {
         debug: true,
+        stream: false,
       },
-      functions: ["CurrentDateTime", "DaysBetweenDates"],
+      agents: ["MathAgent"],
     },
   ],
 };
 
-// Create crew instance with the direct configuration
-const crew = new AxCrew(config, AxCrewFunctions);
-crew.addAgentsToCrew(["Planner", "Calculator", "WebSearch", "Manager"]);
-
 const calculateAndPrintAgentCost = (agent, agentName) => {
   const { models, modelUsage, modelInfo } = agent.axai;
 
@@ -89,16 +97,32 @@ 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 = "best sushi in mumbai";
-console.log(`\n\nQuestion: ${userQuery}`);
+// Create crew instance with the direct configuration
+const crew = new AxCrew(config, AxCrewFunctions);
+crew.addAgentsToCrew(["MathAgent", "ManagerAgent"]);
+
+// const userQuery = "who is considered as the father of the iphone and what is the 7th root of their year of birth";
+const userQuery = "write an essay about the future of ai and then summarize it";
+console.log(`\nTask: ${userQuery}`);
+
+// const Planner = crew.agents.get("Planner");
+// const plannerResponse = await Planner.forward({ task: userQuery });
+// const plan = plannerResponse.plan;
+// console.log(`\n\nPlan: ${plan}`);
+
+const ManagerAgent = crew.agents.get("ManagerAgent");
+const managerAgentResponse = await ManagerAgent.forward({
+  task: userQuery,
+});
+console.log(`\n\nManagerAgent Response: ${managerAgentResponse.result}`);
 
 // const Planner = crew.agents.get("Planner");
 // const plannerResponse = await Planner.forward({ task: userQuery });
 // console.log(`\n\nPlanner Response: ${plannerResponse.plan}`);
 
-const WebSearch = crew.agents.get("WebSearch");
-const webSearchResponse = await WebSearch.forward({ webSearchQuery: userQuery });
-console.log(`\n\nWeb Search Response: ${webSearchResponse.webSearchResponse}`);
+// 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/test1.ts

@@ -0,0 +1,67 @@
+import { AxCrew, AxCrewFunctions, FunctionRegistryType, StateInstance } from './dist/index.js';
+import type { AxFunction } from '@ax-llm/ax';
+
+// Type-safe configuration
+const 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\"",
+    provider: "google-gemini",
+    providerKeyName: "GEMINI_API_KEY",
+    ai: {
+      model: "gemini-1.5-pro",
+      temperature: 0
+    }
+  }]
+};
+
+// Create custom functions with type safety
+class MyCustomFunction {
+  constructor(private state: Record<string, any>) {}
+  
+  toFunction(): AxFunction {
+    return {
+      name: 'MyCustomFunction',
+      description: 'Does something useful',
+      parameters: {
+        type: 'object',
+        properties: {
+          input: { type: 'string', description: "input to the function" }
+        }
+      },
+      func: async ({ input }) => {
+        // Implementation
+        return input;
+      }
+    };
+  }
+}
+
+// Type-safe function registry
+const myFunctions: FunctionRegistryType = {
+  MyCustomFunction
+};
+
+// Create crew with type checking
+const crew = new AxCrew(config, myFunctions);
+
+// Type-safe state management
+crew.state.set('key', 'value');
+const value: string = crew.state.get('key');
+
+// Type-safe agent management
+const agents = crew.addAgentsToCrew(['Planner']);
+const planner = agents?.get('Planner');
+
+if (planner) {
+  // Type-safe agent usage
+  const response = await planner.forward({ task: "Plan something" });
+  console.log(response.plan);
+  const cost = planner.getUsageCost();
+  
+  if (cost) {
+    console.log(`Total cost: $${cost.totalCost}`);
+    console.log(`Total tokens: ${cost.tokenMetrics.totalTokens}`);
+  }
+}