Steps APIs

interface Step<
  ReturnValues extends BaseReturnValues<Record<string, unknown>>,
  YieldValues,
  StepParams,
  NextValues extends BaseNextValue | undefined,
  State extends StateReducer<Record<string, unknown>> = StateReducer<
    Record<string, unknown>
  >,
> {
  /**
   * The name of the executor
   */
  name: string;
  
  /**
   * The beforeExecute function. Used to prepare the state for the Step.
   */
  beforeExecute: (params: StepContext<State>) => Promise<StepParams>;
  
  /**
   * The execute function. Used to execute the Step.
   */
  execute: (
    params: StepParams,
    context: StepContext<State>
  ) => AsyncGenerator<YieldValues, ReturnValues, NextValues>;
  
  /**
   * The afterExecute function. Used to update the state after the Step.
   */
  afterExecute: (
    params: ReturnValues,
    context: StepContext<State>
  ) => Promise<ReturnValues>;
}

interface BaseReturnValues<
  T extends Record<string, unknown> = Record<string, unknown>,
> {
  next: DefaultStepType | (string & {});
  state: T;
  tokenUsage?: AgentTokenUsage;
}

type DefaultStepType =
  | 'agenite.llm-call'
  | 'agenite.tool-call'
  | 'agenite.agent-call'
  | 'agenite.tool-result'
  | 'agenite.end';

interface StepContext<
  Reducer extends StateReducer<Record<string, unknown>>,
> {
  state: StateFromReducer<Reducer>;
  context: Record<string, unknown>;
  agent: Agent<Reducer, any, any, any>;
  parentExecution?: StepContext<any>;
  isNestedExecution?: boolean;
  provider: LLMProvider;
  instructions: string;
  stream: boolean;
  tokenUsage: AgentTokenUsage;
}

// Built-in LLM step implementation
export const LLMStep: Step<
  BaseReturnValues<{
    messages: BaseMessage[];
  }>,
  LLMCallYieldValues,
  LLMCallParams,
  undefined
> = {
  name: 'agenite.llm-call',
  beforeExecute: async (params) => {
    const tools = params.agent.agentConfig.tools || [];
    const agents =
      params.agent.agentConfig.agents?.map((agent) => {
        return {
          name: agent.agentConfig.name,
          description: agent.agentConfig.description || '',
        };
      }) || [];
    return {
      provider: params.provider,
      messages: params.state.messages,
      instructions: params.instructions,
      tools: transformToToolDefinitions([...tools, ...agents]),
      stream: params.stream,
    };
  },
  execute: async function* (params: LLMCallParams) {
    // Call LLM provider and yield streaming responses
    // ...
  },
  afterExecute: async (params) => {
    return params;
  },
};

export const ToolStep: Step<
  BaseReturnValues<{
    messages: BaseMessage[];
  }>,
  ToolCallYieldValues,
  ToolCallParams,
  BaseNextValue
> = {
  name: 'agenite.tool-call',
  beforeExecute: async (params) => {
    // Extract tool call information from messages
    // ...
  },
  execute: async function* (params, context) {
    // Execute the tool and process results
    // ...
  },
  afterExecute: async (params) => {
    return params;
  },
};

export const AgentStep: Step<
  BaseReturnValues<Record<string, unknown>>,
  | AgentCallYieldValues
  | LLMCallYieldValues
  | ToolCallYieldValues
  | YieldAgeniteEnd
  | YieldAgeniteStart
  | ToolResultYieldValues,
  AgentCallParams,
  BaseNextValue
> = {
  name: 'agenite.agent-call',
  beforeExecute: async (params) => {
    // Extract agent call information
    // ...
  },
  execute: async function* (params, context) {
    // Execute nested agent and process results
    // ...
  },
  afterExecute: async (params) => {
    return params;
  },
};

export const ToolResultStep: Step<
  BaseReturnValues<{
    messages: BaseMessage[];
  }>,
  ToolResultYieldValues,
  ToolResultParams,
  BaseNextValue
> = {
  name: 'agenite.tool-result',
  beforeExecute: async (params) => {
    // Prepare tool results 
    // ...
  },
  execute: async function* (params) {
    // Process tool results
    // ...
  },
  afterExecute: async (params) => {
    return params;
  },
};

// Example custom step for data validation
const ValidationStep: Step<
  BaseReturnValues<{
    isValid: boolean;
    errors?: string[];
  }>,
  { type: 'validation.checking'; field: string },
  { data: Record<string, unknown>; rules: ValidationRule[] },
  undefined
> = {
  name: 'validation',
  beforeExecute: async (context) => {
    return {
      data: context.state.formData,
      rules: context.state.validationRules,
    };
  },
  execute: async function* (params, context) {
    const { data, rules } = params;
    const errors: string[] = [];
    
    for (const rule of rules) {
      yield { type: 'validation.checking', field: rule.field };
      
      if (!rule.validate(data[rule.field])) {
        errors.push(`${rule.field}: ${rule.message}`);
      }
    }
    
    return {
      next: errors.length ? 'validation-failed' : 'validation-passed',
      state: {
        isValid: errors.length === 0,
        errors: errors.length ? errors : undefined,
      },
    };
  },
  afterExecute: async (result) => {
    // Perform any cleanup or logging
    return result;
  },
};

// Example custom step for data transformation
const TransformStep: Step<
  BaseReturnValues<{
    transformedData: unknown;
  }>,
  { type: 'transform.processing' },
  { data: unknown; transformFn: (data: unknown) => unknown },
  undefined
> = {
  name: 'transform',
  beforeExecute: async (context) => {
    return {
      data: context.state.rawData,
      transformFn: context.state.transformFunction,
    };
  },
  execute: async function* ({ data, transformFn }) {
    yield { type: 'transform.processing' };
    
    const transformedData = transformFn(data);
    
    return {
      next: 'agenite.llm-call',
      state: {
        transformedData,
      },
    };
  },
  afterExecute: async (result) => {
    return result;
  },
};

const agent = new Agent({
  name: 'custom-steps-agent',
  provider: new BedrockProvider({ model: 'anthropic.claude-3-5-sonnet-20240620-v1:0' }),
  steps: {
    // Include custom steps
    'validation': ValidationStep,
    'transform': TransformStep,
    // Include default steps
    'agenite.llm-call': LLMStep,
    'agenite.tool-call': ToolStep,
    'agenite.agent-call': AgentStep,
    'agenite.tool-result': ToolResultStep,
  },
  // Specify the starting step
  startStep: 'validation',
});