diff options
Diffstat (limited to 'src/client/views/nodes/chatbot/tools/DataAnalysisTool.ts')
-rw-r--r-- | src/client/views/nodes/chatbot/tools/DataAnalysisTool.ts | 67 |
1 files changed, 67 insertions, 0 deletions
diff --git a/src/client/views/nodes/chatbot/tools/DataAnalysisTool.ts b/src/client/views/nodes/chatbot/tools/DataAnalysisTool.ts new file mode 100644 index 000000000..8c5e3d9cd --- /dev/null +++ b/src/client/views/nodes/chatbot/tools/DataAnalysisTool.ts @@ -0,0 +1,67 @@ +import { Observation } from '../types/types'; +import { ParametersType, ToolInfo } from '../types/tool_types'; +import { BaseTool } from './BaseTool'; + +const dataAnalysisToolParams = [ + { + name: 'csv_file_names', + type: 'string[]', + description: 'List of names of the CSV files to analyze', + required: true, + max_inputs: 3, + }, +] as const; + +type DataAnalysisToolParamsType = typeof dataAnalysisToolParams; + +const dataAnalysisToolInfo: ToolInfo<DataAnalysisToolParamsType> = { + name: 'dataAnalysis', + description: 'Provides the full CSV file text for your analysis based on the user query and the available CSV file(s).', + citationRules: 'No citation needed.', + parameterRules: dataAnalysisToolParams, +}; + +export class DataAnalysisTool extends BaseTool<DataAnalysisToolParamsType> { + private csv_files_function: () => { filename: string; id: string; text: string }[]; + + constructor(csv_files: () => { filename: string; id: string; text: string }[]) { + super(dataAnalysisToolInfo); + this.csv_files_function = csv_files; + } + + getFileContent(filename: string): string | undefined { + const files = this.csv_files_function(); + const file = files.find(f => f.filename === filename); + return file?.text; + } + + getFileID(filename: string): string | undefined { + const files = this.csv_files_function(); + const file = files.find(f => f.filename === filename); + return file?.id; + } + + async execute(args: ParametersType<DataAnalysisToolParamsType>): Promise<Observation[]> { + const filenames = args.csv_file_names; + const results: Observation[] = []; + + for (const filename of filenames) { + const fileContent = this.getFileContent(filename); + const fileID = this.getFileID(filename); + + if (fileContent && fileID) { + results.push({ + type: 'text', + text: `<chunk chunk_id="${fileID}" chunk_type="csv">${fileContent}</chunk>`, + }); + } else { + results.push({ + type: 'text', + text: `File not found: ${filename}`, + }); + } + } + + return results; + } +} |