making it work even better

1 files changed, 84 insertions, 62 deletions

diff --git a/src/client/views/nodes/chatbot/tools/CreateAnyDocTool.ts b/src/client/views/nodes/chatbot/tools/CreateAnyDocTool.ts
index bb1761cee..6f61b77d4 100644
--- a/src/client/views/nodes/chatbot/tools/CreateAnyDocTool.ts
+++ b/src/client/views/nodes/chatbot/tools/CreateAnyDocTool.ts
@@ -1,38 +1,51 @@
 import { v4 as uuidv4 } from 'uuid';
 import { BaseTool } from './BaseTool';
 import { Observation } from '../types/types';
-import { ParametersType, TypeMap, Parameter } from '../types/tool_types';
+import { ParametersType, Parameter } from '../types/tool_types';
 import { DocumentOptions, Docs } from '../../../../documents/Documents';
 
 /**
- * List of supported document types.
+ * List of supported document types that can be created via text LLM.
  */
-const supportedDocumentTypes = [
-    'text',
-    'image',
-    'pdf',
-    'video',
-    'audio',
-    'web',
-    'map',
-    'equation',
-    'functionPlot',
-    'dataViz',
-    'chat',
-    // Add more document types as needed
-];
+type supportedDocumentTypesType = 'text' | 'html' | 'equation' | 'functionPlot' | 'dataviz' | 'noteTaking' | 'rtf' | 'message';
+const supportedDocumentTypes: supportedDocumentTypesType[] = ['text', 'html', 'equation', 'functionPlot', 'dataviz', 'noteTaking', 'rtf', 'message'];
 
 /**
- * Description of document options for each type.
+ * Description of document options and data field for each type.
  */
-const documentOptionsDescription = {
-    text: ['title', 'backgroundColor', 'fontColor', 'text_align', 'layout', 'text_content'],
-    image: ['title', 'backgroundColor', 'width', 'height', 'layout'],
-    pdf: ['title', 'backgroundColor', 'width', 'height', 'layout'],
-    video: ['title', 'backgroundColor', 'width', 'height', 'layout'],
-    audio: ['title', 'backgroundColor', 'layout'],
-    web: ['title', 'backgroundColor', 'width', 'height', 'layout', 'url'],
-    // Include descriptions for other document types
+const documentTypesInfo = {
+    text: {
+        options: ['title', 'backgroundColor', 'fontColor', 'text_align', 'layout'],
+        dataDescription: 'The text content of the document.',
+    },
+    html: {
+        options: ['title', 'backgroundColor', 'layout'],
+        dataDescription: 'The HTML-formatted text content of the document.',
+    },
+    equation: {
+        options: ['title', 'backgroundColor', 'fontColor', 'layout'],
+        dataDescription: 'The equation content as a string.',
+    },
+    functionPlot: {
+        options: ['title', 'backgroundColor', 'layout', 'function_definition'],
+        dataDescription: 'The function definition(s) for plotting. Provide as a string or array of function definitions.',
+    },
+    dataviz: {
+        options: ['title', 'backgroundColor', 'layout', 'chartType'],
+        dataDescription: 'A string of comma-separated values representing the CSV data.',
+    },
+    noteTaking: {
+        options: ['title', 'backgroundColor', 'layout'],
+        dataDescription: 'The initial content or structure for note-taking.',
+    },
+    rtf: {
+        options: ['title', 'backgroundColor', 'layout'],
+        dataDescription: 'The rich text content in RTF format.',
+    },
+    message: {
+        options: ['title', 'backgroundColor', 'layout'],
+        dataDescription: 'The message content of the document.',
+    },
 };
 
 const createAnyDocumentToolParams = [
@@ -45,19 +58,19 @@ const createAnyDocumentToolParams = [
     {
         name: 'data',
         type: 'string',
-        description: 'The content or data of the document (e.g., text content, URL, etc.).',
-        required: false,
+        description: 'The content or data of the document. The exact format depends on the document type.',
+        required: true,
     },
     {
         name: 'options',
         type: 'string',
-        description: `A JSON string representing the document options. Available options depend on the document type.\n
-        For example, for 'text' documents, options include: ${documentOptionsDescription['text'].join(', ')}.\n
-        For 'image' documents, options include: ${documentOptionsDescription['image'].join(', ')}.\n
-        For 'pdf' documents, options include: ${documentOptionsDescription['pdf'].join(', ')}.\n
-        For 'video' documents, options include: ${documentOptionsDescription['video'].join(', ')}.\n
-        For 'audio' documents, options include: ${documentOptionsDescription['audio'].join(', ')}.\n
-        For 'web' documents, options include: ${documentOptionsDescription['web'].join(', ')}.\n`,
+        description: `A JSON string representing the document options. Available options depend on the document type. For example:
+${supportedDocumentTypes
+    .map(
+        docType => `
+- For '${docType}' documents, options include: ${documentTypesInfo[docType].options.join(', ')}`
+    )
+    .join('\n')}`,
         required: false,
     },
 ] as const;
@@ -70,23 +83,41 @@ export class CreateAnyDocumentTool extends BaseTool<CreateAnyDocumentToolParamsT
     constructor(addLinkedDoc: (doc_type: string, data: string | undefined, options: DocumentOptions, id: string) => void) {
         super(
             'createAnyDocument',
-            `Creates any type of document with the provided options and data. Supported document types are: ${supportedDocumentTypes.join(', ')}.`,
+            `Creates any type of document with the provided options and data. Supported document types are: ${supportedDocumentTypes.join(', ')}. dataviz is a csv table tool, so for CSVs, use dataviz. Here are the options for each type:
+            <supported_document_types>
+        ${supportedDocumentTypes
+            .map(
+                docType => `
+        <document_type name="${docType}">
+          <data_description>${documentTypesInfo[docType].dataDescription}</data_description>
+          <options>
+            ${documentTypesInfo[docType].options.map(option => `<option>${option}</option>`).join('\n')}
+          </options>
+        </document_type>
+        `
+            )
+            .join('\n')}
+      </supported_document_types>`,
             createAnyDocumentToolParams,
             'Provide the document type, data, and options for the document. Options should be a valid JSON string containing the document options specific to the document type.',
-            'Creates any type of document with the provided options and data.'
+            `Creates any type of document with the provided options and data. Supported document types are: ${supportedDocumentTypes.join(', ')}.`
         );
         this._addLinkedDoc = addLinkedDoc;
     }
 
     async execute(args: ParametersType<CreateAnyDocumentToolParamsType>): Promise<Observation[]> {
         try {
-            const documentType = args.document_type.toLowerCase();
+            const documentType: supportedDocumentTypesType = args.document_type.toLowerCase() as supportedDocumentTypesType;
             let options: DocumentOptions = {};
 
             if (!supportedDocumentTypes.includes(documentType)) {
                 throw new Error(`Unsupported document type: ${documentType}. Supported types are: ${supportedDocumentTypes.join(', ')}.`);
             }
 
+            if (!args.data) {
+                throw new Error(`Data is required for ${documentType} documents. ${documentTypesInfo[documentType].dataDescription}`);
+            }
+
             if (args.options) {
                 try {
                     options = JSON.parse(args.options as string) as DocumentOptions;
@@ -95,37 +126,28 @@ export class CreateAnyDocumentTool extends BaseTool<CreateAnyDocumentToolParamsT
                 }
             }
 
-            const data = args.data as string | undefined;
+            const data = args.data as string;
             const id = uuidv4();
 
-            // Validate and set default options based on document type
-            switch (documentType) {
-                case 'text':
-                    if (!data) {
-                        throw new Error('Data is required for text documents.');
-                    }
-                    options.title = options.title || 'New Text Document';
-                    break;
-                case 'image':
-                case 'pdf':
-                case 'video':
-                case 'audio':
-                case 'web':
-                    if (!data) {
-                        throw new Error(`Data (e.g., URL) is required for ${documentType} documents.`);
-                    }
-                    options.title = options.title || `New ${documentType.charAt(0).toUpperCase() + documentType.slice(1)} Document`;
-                    break;
-                // Add cases and default options for other document types as needed
-                default:
-                    break;
-            }
+            // Set default options if not provided
+            options.title = options.title || `New ${documentType.charAt(0).toUpperCase() + documentType.slice(1)} Document`;
 
+            // Call the function to add the linked document
             this._addLinkedDoc(documentType, data, options, id);
 
-            return [{ type: 'text', text: `Created ${documentType} document with ID ${id}.` }];
+            return [
+                {
+                    type: 'text',
+                    text: `Created ${documentType} document with ID ${id}.`,
+                },
+            ];
         } catch (error) {
-            return [{ type: 'text', text: 'Error creating document: ' + (error as Error).message }];
+            return [
+                {
+                    type: 'text',
+                    text: 'Error creating document: ' + (error as Error).message,
+                },
+            ];
         }
     }
 }