aboutsummaryrefslogtreecommitdiff
path: root/src/client/views/nodes/chatbot/tools/DataAnalysisTool.ts
blob: 2e663fed1cb164d09a542e47f480fa47100432e9 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
import { BaseTool } from './BaseTool';

export class DataAnalysisTool extends BaseTool<{ csv_file_name: string | string[] }> {
    private csv_files_function: () => { filename: string; id: string; text: string }[];

    constructor(csv_files: () => { filename: string; id: string; text: string }[]) {
        super(
            'dataAnalysis',
            'Analyzes, and provides insights, from one or more CSV files',
            {
                csv_file_name: {
                    type: 'string',
                    description: 'Name(s) of the CSV file(s) to analyze',
                    required: 'true',
                    max_inputs: '3',
                },
            },
            'Provide the name(s) of up to 3 CSV files to analyze based on the user query and whichever available CSV files may be relevant.',
            'Provides the full CSV file text for your analysis based on the user query and the available CSV file(s). '
        );
        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: { csv_file_name: string | string[] }): Promise<unknown> {
        const filenames = Array.isArray(args.csv_file_name) ? args.csv_file_name : [args.csv_file_name];
        const results = [];

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