aboutsummaryrefslogtreecommitdiff
path: root/src/client/views/nodes/chatbot/tools/DataAnalysisTool.ts
diff options
context:
space:
mode:
Diffstat (limited to 'src/client/views/nodes/chatbot/tools/DataAnalysisTool.ts')
-rw-r--r--src/client/views/nodes/chatbot/tools/DataAnalysisTool.ts67
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;
+ }
+}