@amitdeshmukh/ax-crew 3.3.2 → 3.3.3

package/CHANGELOG.md

@@ -5,6 +5,18 @@ This Changelog format is based on [Keep a Changelog]
 adheres to [Semantic Versioning](https://semver.org/spec/
 v2.0.0.html).
 
+## [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

@@ -22,6 +22,86 @@ AxLLM is a peer dependency, so you will need to install it separately.
 npm install @ax-llm/ax@latest
 ```
 
+### 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/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/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.3",
   "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",

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 })

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.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}`);
+  }
+}