1 files changed, 55 insertions, 23 deletions

diff --git a/src/client/views/nodes/chatbot/tools/CreateAnyDocTool.ts b/src/client/views/nodes/chatbot/tools/CreateAnyDocTool.ts
index 7ec39704a..5cf858998 100644
--- a/src/client/views/nodes/chatbot/tools/CreateAnyDocTool.ts
+++ b/src/client/views/nodes/chatbot/tools/CreateAnyDocTool.ts
@@ -1,10 +1,12 @@
 import { toLower } from 'lodash';
+import { Doc } from '../../../../../fields/Doc';
+import { Id } from '../../../../../fields/FieldSymbols';
 import { DocumentOptions } from '../../../../documents/Documents';
-import { ParametersType } from '../types/tool_types';
+import { parsedDoc } from '../chatboxcomponents/ChatBox';
+import { ParametersType, ToolInfo } from '../types/tool_types';
 import { Observation } from '../types/types';
 import { BaseTool } from './BaseTool';
 import { supportedDocumentTypes } from './CreateDocumentTool';
-import { Id } from '../../../../../fields/FieldSymbols';
 
 const standardOptions = ['title', 'backgroundColor'];
 /**
@@ -43,10 +45,34 @@ const documentTypesInfo: { [key in supportedDocumentTypes]: { options: string[];
         options: standardOptions,
         dataDescription: 'The rich text content in RTF format.',
     },
+    [supportedDocumentTypes.image]: {
+        options: standardOptions,
+        dataDescription: 'The image content as an image file URL.',
+    },
+    [supportedDocumentTypes.pdf]: {
+        options: standardOptions,
+        dataDescription: 'the pdf content as a PDF file url.',
+    },
+    [supportedDocumentTypes.audio]: {
+        options: standardOptions,
+        dataDescription: 'The audio content as a file url.',
+    },
+    [supportedDocumentTypes.video]: {
+        options: standardOptions,
+        dataDescription: 'The video content as a file url.',
+    },
     [supportedDocumentTypes.message]: {
         options: standardOptions,
         dataDescription: 'The message content of the document.',
     },
+    [supportedDocumentTypes.mermaid]: {
+        options: ['title', 'backgroundColor'],
+        dataDescription: 'The Mermaid diagram content.',
+    },
+    [supportedDocumentTypes.script]: {
+        options: ['title', 'backgroundColor'],
+        dataDescription: 'The compilable JavaScript code. Use this for creating scripts.',
+    },
 };
 
 const createAnyDocumentToolParams = [
@@ -75,28 +101,34 @@ const createAnyDocumentToolParams = [
 
 type CreateAnyDocumentToolParamsType = typeof createAnyDocumentToolParams;
 
+const createAnyDocToolInfo: ToolInfo<CreateAnyDocumentToolParamsType> = {
+    name: 'createAnyDocument',
+    description:
+        `Creates any type of document with the provided options and data. 
+      Supported document types are: ${Object.values(supportedDocumentTypes).join(', ')}. 
+           dataviz is a csv table tool, so for CSVs, use dataviz. Here are the options for each type:
+        <supported_document_types>` +
+        Object.entries(documentTypesInfo)
+            .map(
+                ([doc_type, info]) =>
+                    `<document_type name="${doc_type}">
+                    <data_description>${info.dataDescription}</data_description>
+                    <options>` +
+                    info.options.map(option => `<option>${option}</option>`).join('\n') +
+                    `</options>
+                </document_type>`
+            )
+            .join('\n') +
+        `</supported_document_types>`,
+    parameterRules: createAnyDocumentToolParams,
+    citationRules: 'No citation needed.',
+};
+
 export class CreateAnyDocumentTool extends BaseTool<CreateAnyDocumentToolParamsType> {
-    private _addLinkedDoc: (doc_type: supportedDocumentTypes, data: unknown, options: DocumentOptions) => void;
+    private _addLinkedDoc: (doc: parsedDoc) => Doc | undefined;
 
-    constructor(addLinkedDoc: (doc_type: supportedDocumentTypes, data: unknown, options: DocumentOptions) => void) {
-        // prettier-ignore
-        super(
-            'createAnyDocument',
-            `Creates any type of document with the provided options and data. Supported document types are: ${Object.values(supportedDocumentTypes).join(', ')}. dataviz is a csv table tool, so for CSVs, use dataviz. Here are the options for each type:
-            <supported_document_types>
-                ${Object.entries(documentTypesInfo).map(([doc_type, info]) => `
-                <document_type name="${doc_type}">
-                    <data_description>${info.dataDescription}</data_description>
-                    <options>
-                        ${info.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. Supported document types are: ${Object.values(supportedDocumentTypes).join(', ')}.`
-        );
+    constructor(addLinkedDoc: (doc: parsedDoc) => Doc | undefined) {
+        super(createAnyDocToolInfo);
         this._addLinkedDoc = addLinkedDoc;
     }
 
@@ -116,7 +148,7 @@ export class CreateAnyDocumentTool extends BaseTool<CreateAnyDocumentToolParamsT
             const options: DocumentOptions = !args.options ? {} : JSON.parse(args.options);
 
             // Call the function to add the linked document (add default title that can be overriden if set in options)
-            const doc = this._addLinkedDoc(documentType, args.data, { title: `New ${documentType.charAt(0).toUpperCase() + documentType.slice(1)} Document`, ...options });
+            const doc = this._addLinkedDoc({ doc_type: documentType, data: args.data, title: `New ${documentType.charAt(0).toUpperCase() + documentType.slice(1)} Document`, ...options });
 
             return [{ type: 'text', text: `Created ${documentType} document with ID ${doc?.[Id]}.` }];
         } catch (error) {