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

A cool library focusing on handling various flows

import { BaseAINode, BaseAINodeInput } from './baseAINode';
import { INodeContext } from '@flowlab/core';
import { IAIProvider, AIModelOptions } from '../types';

interface StructuredDataExtractionNodeInput extends BaseAINodeInput {
  sourceText?: string; // Text to extract from (can be templated)
  sourceTextTemplate?: string; // Alternative to sourceText for templating
  jsonSchema: object; // The JSON schema for extraction
}

export class StructuredDataExtractionNode extends BaseAINode {
  id = 'StructuredDataExtractionNode';

  // MARK: 执行AI
  protected async executeAI(context: INodeContext, provider: IAIProvider, renderedInput: any): Promise<Record<string, any>> {
     const nodeInput = context.input as StructuredDataExtractionNodeInput;
     const options = nodeInput.modelOptions;
     let textToProcess: string;

     if (nodeInput.sourceTextTemplate) {
         textToProcess = this.renderPromptTemplate(nodeInput.sourceTextTemplate, nodeInput.inputVariables || {}, context);
     } else if (nodeInput.sourceText) {
         textToProcess = nodeInput.sourceText;
     } else {
         throw new Error("Missing 'sourceText' or 'sourceTextTemplate' input for StructuredDataExtractionNode");
     }

     if (!nodeInput.jsonSchema || typeof nodeInput.jsonSchema !== 'object') {
         throw new Error("Missing or invalid 'jsonSchema' object input for StructuredDataExtractionNode");
     }

     const { data, usage } = await provider.extractStructuredData(textToProcess, nodeInput.jsonSchema, options, context);
     return { extractedData: data, usage };
  }

  // MARK: 验证输入
  validateInput(input: Record<string, any>): boolean {
      if (!super.validateInput(input)) return false;
      if (!input.jsonSchema || typeof input.jsonSchema !== 'object') {
        console.error(`Validation Error [${this.id}]: Missing or invalid 'jsonSchema'.`);
        return false;
      }
       if (!input.sourceText && !input.sourceTextTemplate) {
         console.error(`Validation Error [${this.id}]: Missing 'sourceText' or 'sourceTextTemplate'.`);
         return false;
       }
      return true;
    }
}