@flowlab/all/packages/ai/src/nodes/baseAINode.ts

A cool library focusing on handling various flows

import {
  BaseNode, INodeContext, INodeOutput, NodeStatus, IWorkflowContext,
  IScheduler, // Import Scheduler interface from core
  WorkflowError
} from '@flowlab/core';
import { AIProviderRegistry, IAIProvider, AIModelOptions } from '../types';
import { aiProviderRegistry } from '../providerRegistry'; // Use the singleton registry


// MARK: 基础AI节点输入
// Define standard input properties expected by AI nodes
export interface BaseAINodeInput {
  providerName: string;         // Name of the AI provider (e.g., 'openai')
  promptTemplate?: string;        // Template string for prompts
  modelOptions?: AIModelOptions;  // Overrides for model parameters
  providerConfig?: object;      // Override provider config for this node
  inputVariables?: Record<string, string>; // Explicit mapping from context vars to template vars
  scheduleTask?: boolean;       // Flag to request scheduling via core scheduler
  scheduleOptions?: any;        // Options for the scheduler (delay, priority etc.)
}

// MARK: 基础AI节点
export abstract class BaseAINode extends BaseNode {
  // Access the shared registry
  protected registry: AIProviderRegistry = aiProviderRegistry;
  // Scheduler instance, potentially injected via executor options
  protected scheduler?: IScheduler;

  constructor(scheduler?: IScheduler) { // Accept scheduler optionally
      super();
      this.scheduler = scheduler;
  }

  // MARK: 执行AI
  // Abstract method for the core AI logic
  protected abstract executeAI(context: INodeContext, provider: IAIProvider, renderedInput: any): Promise<Record<string, any>>;

  // MARK: 执行
  async execute(context: INodeContext): Promise<INodeOutput> {
    const nodeInput = context.input as BaseAINodeInput;

    // --- Scheduling Logic ---
    if (nodeInput.scheduleTask === true) {
      if (!this.scheduler) {
        context.logs.push('Warning: scheduleTask=true but no scheduler configured/available.');
      } else {
        try {
          // NOTE: Scheduling requires context serialization and a mechanism
          // for the worker to report back. This is a complex topic involving
          // potentially durable workflows and state persistence.
          // Here, we just initiate the scheduling. FlowLab core needs to handle the rest.
          const taskId = await this.scheduler.scheduleNode(
            this.id, // Schedule this node type
            { // Pass necessary context parts (needs careful selection/serialization)
              ...context,
              input: context.input, // Ensure input is passed
              workflowContext: { // Only pass essential workflow context if needed
                workflowId: context.workflowContext.workflowId,
                variables: context.workflowContext.variables,
                // Potentially tenantId, userId if needed for scheduled task logic
              },
              // Remove potentially non-serializable parts or large objects if needed
            },
            nodeInput.scheduleOptions || {}
          );
          // Indicate the task is scheduled and waiting for the external worker
          return {
            status: NodeStatus.PENDING, // Or a custom 'SCHEDULED' status if core supports it
            output: { scheduledTaskId: taskId }
          };
        } catch (error: any) {
            context.logs.push(`Error scheduling AI task: ${error.message}`);
            context.error = new WorkflowError(`Failed to schedule AI task: ${error.message}`);
            return { status: NodeStatus.FAILED, output: { error: context.error.message } };
        }
      }
    }

    // MARK: 直接执行
    try {
      const provider = this.registry.getProvider(nodeInput.providerName, nodeInput.providerConfig);

      // Render input (e.g., prompt template)
      const renderedInput = this.renderInput(nodeInput, context);

      // Delegate to specific AI execution
      const aiResult = await this.executeAI(context, provider, renderedInput);

      // Standardize output structure slightly
      return {
        status: NodeStatus.COMPLETED,
        output: {
          ...aiResult, // Include AI-specific results (text, data, usage etc.)
          provider: nodeInput.providerName
        }
      };

    } catch (error: any) {
      context.logs.push(`AI Node execution failed: ${error.message}`);
      context.error = error instanceof Error ? error : new NodeExecutionError(String(error), this.id, context.stepId);
      // TODO: Map specific AI provider errors (e.g., rate limit, content filter) if possible
      return {
        status: NodeStatus.FAILED,
        output: { error: context.error.message, provider: nodeInput.providerName }
      };
    }
  }

  // MARK: 输入渲染（e.g., Prompt Templating）
  protected renderInput(nodeInput: BaseAINodeInput, context: INodeContext): any {
    if (nodeInput.promptTemplate) {
      return this.renderPromptTemplate(nodeInput.promptTemplate, nodeInput.inputVariables || {}, context);
    }
    // If no template, pass other relevant inputs directly (needs refinement based on node type)
    return { ...nodeInput };
  }

  // MARK: 简单模板引擎（replace with a more robust one if needed）
  protected renderPromptTemplate(template: string, mapping: Record<string, string>, context: INodeContext): string {
    let rendered = template;
    const combinedContext = {
      ...context.workflowContext.variables, // workflow variables
      ...context.input // Step input (use mapping if provided)
    };

    // Priority 1: Explicit mapping from inputVariables
    for (const templateVar in mapping) {
        const contextPath = mapping[templateVar]; // e.g., "workflow.variables.someVar" or "input.someInput"
        const value = this.resolveContextPath(contextPath, context);
        const regex = new RegExp(`\\$\\{${templateVar}\\}`, 'g');
        rendered = rendered.replace(regex, String(value ?? ''));
    }

    // MARK: 优先2：直接上下文访问（简单变量）
    // Matches ${workflow.variables.varName} or ${input.varName}
    rendered = rendered.replace(/\$\{(workflow\.variables|input)\.(\w+)\}/g, (match, scope, varName) => {
        const contextPath = `${scope}.${varName}`;
        if (!mapping || !Object.values(mapping).includes(contextPath)) { // Avoid double replacement
           const value = this.resolveContextPath(contextPath, context);
           return String(value ?? '');
        }
        return match; // Already replaced by mapping
    });

    // Log the rendered prompt for debugging (optional)
    context.logs.push(`Rendered Prompt: ${rendered.substring(0, 100)}...`);
    return rendered;
  }

  // MARK: 解析简单点符号路径
  private resolveContextPath(path: string, context: INodeContext): any {
      const parts = path.split('.');
      let currentVal: any = null;
      if (parts[0] === 'workflow' && parts[1] === 'variables') {
          currentVal = context.workflowContext.variables;
          parts.splice(0, 2);
      } else if (parts[0] === 'input') {
          currentVal = context.input;
          parts.splice(0, 1);
      } else {
          return undefined; // Unsupported path root
      }

      for (const part of parts) {
          if (currentVal === null || typeof currentVal !== 'object') return undefined;
          currentVal = currentVal[part];
      }
      return currentVal;
  }

  // MARK: 默认输入验证（在特定节点中重写）
  // Default input validation (override in specific nodes)
  validateInput(input: Record<string, any>): boolean {
    if (!input.providerName || typeof input.providerName !== 'string') {
      console.error(`Validation Error [${this.id}]: Missing or invalid 'providerName'.`);
      return false;
    }
    // Specific nodes will add checks for promptTemplate, schema, labels etc.
    return true;
  }

  // Optional: 如果AI操作需要回滚，则实现补偿

  // async compensate(context: INodeContext): Promise<void> { ... }
}