1 files changed, 158 insertions, 0 deletions
diff --git a/src/client/views/nodes/chatbot/tools/CreateAnyDocTool.ts b/src/client/views/nodes/chatbot/tools/CreateAnyDocTool.ts
new file mode 100644
index 000000000..754d230c8
--- /dev/null
+++ b/src/client/views/nodes/chatbot/tools/CreateAnyDocTool.ts
@@ -0,0 +1,158 @@
+import { toLower } from 'lodash';
+import { Doc } from '../../../../../fields/Doc';
+import { Id } from '../../../../../fields/FieldSymbols';
+import { DocumentOptions } from '../../../../documents/Documents';
+import { parsedDoc } from '../chatboxcomponents/ChatBox';
+import { ParametersType, ToolInfo } from '../types/tool_types';
+import { Observation } from '../types/types';
+import { BaseTool } from './BaseTool';
+import { supportedDocTypes } from './CreateDocumentTool';
+
+const standardOptions = ['title', 'backgroundColor'];
+/**
+ * Description of document options and data field for each type.
+ */
+const documentTypesInfo: { [key in supportedDocTypes]: { options: string[]; dataDescription: string } } = {
+    [supportedDocTypes.flashcard]: {
+        options: [...standardOptions, 'fontColor', 'text_align'],
+        dataDescription: 'an array of two strings.  the first string contains a question, and the second string contains an answer',
+    },
+    [supportedDocTypes.text]: {
+        options: [...standardOptions, 'fontColor', 'text_align'],
+        dataDescription: 'The text content of the document.',
+    },
+    [supportedDocTypes.html]: {
+        options: [],
+        dataDescription: 'The HTML-formatted text content of the document.',
+    },
+    [supportedDocTypes.equation]: {
+        options: [...standardOptions, 'fontColor'],
+        dataDescription: 'The equation content as a string.',
+    },
+    [supportedDocTypes.functionplot]: {
+        options: [...standardOptions, 'function_definition'],
+        dataDescription: 'The function definition(s) for plotting. Provide as a string or array of function definitions.',
+    },
+    [supportedDocTypes.dataviz]: {
+        options: [...standardOptions, 'chartType'],
+        dataDescription: 'A string of comma-separated values representing the CSV data.',
+    },
+    [supportedDocTypes.notetaking]: {
+        options: standardOptions,
+        dataDescription: 'The initial content or structure for note-taking.',
+    },
+    [supportedDocTypes.rtf]: {
+        options: standardOptions,
+        dataDescription: 'The rich text content in RTF format.',
+    },
+    [supportedDocTypes.image]: {
+        options: standardOptions,
+        dataDescription: 'The image content as an image file URL.',
+    },
+    [supportedDocTypes.pdf]: {
+        options: standardOptions,
+        dataDescription: 'the pdf content as a PDF file url.',
+    },
+    [supportedDocTypes.audio]: {
+        options: standardOptions,
+        dataDescription: 'The audio content as a file url.',
+    },
+    [supportedDocTypes.video]: {
+        options: standardOptions,
+        dataDescription: 'The video content as a file url.',
+    },
+    [supportedDocTypes.message]: {
+        options: standardOptions,
+        dataDescription: 'The message content of the document.',
+    },
+    [supportedDocTypes.diagram]: {
+        options: ['title', 'backgroundColor'],
+        dataDescription: 'diagram content as a text string in Mermaid format.',
+    },
+    [supportedDocTypes.script]: {
+        options: ['title', 'backgroundColor'],
+        dataDescription: 'The compilable JavaScript code. Use this for creating scripts.',
+    },
+};
+
+const createAnyDocumentToolParams = [
+    {
+        name: 'document_type',
+        type: 'string',
+        description: `The type of the document to create. Supported types are: ${Object.values(supportedDocTypes).join(', ')}`,
+        required: true,
+    },
+    {
+        name: 'data',
+        type: 'string',
+        description: 'The content or data of the document. The exact format depends on the document type.',
+        required: true,
+    },
+    {
+        name: 'options',
+        type: 'string',
+        required: false,
+        description: `A JSON string representing the document options. Available options depend on the document type. For example:
+            ${Object.entries(documentTypesInfo).map( ([doc_type, info]) => `
+- For '${doc_type}' documents, options include: ${info.options.join(', ')}`) 
+                    .join('\n')}`, // prettier-ignore
+    },
+] as const;
+
+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(supportedDocTypes).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: parsedDoc) => Doc | undefined;
+
+    constructor(addLinkedDoc: (doc: parsedDoc) => Doc | undefined) {
+        super(createAnyDocToolInfo);
+        this._addLinkedDoc = addLinkedDoc;
+    }
+
+    async execute(args: ParametersType<CreateAnyDocumentToolParamsType>): Promise<Observation[]> {
+        try {
+            const documentType = toLower(args.document_type) as unknown as supportedDocTypes;
+            const info = documentTypesInfo[documentType];
+
+            if (info === undefined) {
+                throw new Error(`Unsupported document type: ${documentType}. Supported types are: ${Object.values(supportedDocTypes).join(', ')}.`);
+            }
+
+            if (!args.data) {
+                throw new Error(`Data is required for ${documentType} documents. ${info.dataDescription}`);
+            }
+
+            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({ 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) {
+            return [{ type: 'text', text: 'Error creating document: ' + (error as Error).message }];
+        }
+    }
+}