Merge branch 'DocCreatorMenu-work' of https://github.com/brown-dash/Dash-Web into DocCreatorMenu-work

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