updated from master

1 files changed, 83 insertions, 87 deletions

diff --git a/src/client/views/nodes/chatbot/tools/CreateAnyDocTool.ts b/src/client/views/nodes/chatbot/tools/CreateAnyDocTool.ts
index 5f3af8296..5cf858998 100644
--- a/src/client/views/nodes/chatbot/tools/CreateAnyDocTool.ts
+++ b/src/client/views/nodes/chatbot/tools/CreateAnyDocTool.ts
@@ -1,57 +1,76 @@
-import { v4 as uuidv4 } from 'uuid';
-import { BaseTool } from './BaseTool';
+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 { ParametersType, Parameter, ToolInfo } from '../types/tool_types';
-import { DocumentOptions, Docs } from '../../../../documents/Documents';
-
-/**
- * List of supported document types that can be created via text LLM.
- */
-type supportedDocumentTypesType = 'text' | 'html' | 'equation' | 'function_plot' | 'dataviz' | 'note_taking' | 'rtf' | 'message' | 'mermaid_diagram' | 'script';
-const supportedDocumentTypes: supportedDocumentTypesType[] = ['text', 'html', 'equation', 'function_plot', 'dataviz', 'note_taking', 'rtf', 'message', 'mermaid_diagram', 'script'];
+import { BaseTool } from './BaseTool';
+import { supportedDocumentTypes } from './CreateDocumentTool';
 
+const standardOptions = ['title', 'backgroundColor'];
 /**
  * Description of document options and data field for each type.
  */
-const documentTypesInfo = {
-    text: {
-        options: ['title', 'backgroundColor', 'fontColor', 'text_align', 'layout'],
-        dataDescription: 'The text content of the text document. Should contain all the text content.',
+const documentTypesInfo: { [key in supportedDocumentTypes]: { options: string[]; dataDescription: string } } = {
+    [supportedDocumentTypes.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',
+    },
+    [supportedDocumentTypes.text]: {
+        options: [...standardOptions, 'fontColor', 'text_align'],
+        dataDescription: 'The text content of the document.',
     },
-    html: {
-        options: ['title', 'backgroundColor', 'layout'],
+    [supportedDocumentTypes.html]: {
+        options: [],
         dataDescription: 'The HTML-formatted text content of the document.',
     },
-    equation: {
-        options: ['title', 'backgroundColor', 'fontColor', 'layout'],
+    [supportedDocumentTypes.equation]: {
+        options: [...standardOptions, 'fontColor'],
         dataDescription: 'The equation content as a string.',
     },
-    function_plot: {
-        options: ['title', 'backgroundColor', 'layout', 'function_definition'],
+    [supportedDocumentTypes.functionplot]: {
+        options: [...standardOptions, 'function_definition'],
         dataDescription: 'The function definition(s) for plotting. Provide as a string or array of function definitions.',
     },
-    dataviz: {
-        options: ['title', 'backgroundColor', 'layout', 'chartType'],
+    [supportedDocumentTypes.dataviz]: {
+        options: [...standardOptions, 'chartType'],
         dataDescription: 'A string of comma-separated values representing the CSV data.',
     },
-    note_taking: {
-        options: ['title', 'backgroundColor', 'layout'],
+    [supportedDocumentTypes.notetaking]: {
+        options: standardOptions,
         dataDescription: 'The initial content or structure for note-taking.',
     },
-    rtf: {
-        options: ['title', 'backgroundColor', 'layout'],
+    [supportedDocumentTypes.rtf]: {
+        options: standardOptions,
         dataDescription: 'The rich text content in RTF format.',
     },
-    message: {
-        options: ['title', 'backgroundColor', 'layout'],
+    [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.',
     },
-    mermaid_diagram: {
-        options: ['title', 'backgroundColor', 'layout'],
+    [supportedDocumentTypes.mermaid]: {
+        options: ['title', 'backgroundColor'],
         dataDescription: 'The Mermaid diagram content.',
     },
-    script: {
-        options: ['title', 'backgroundColor', 'layout'],
+    [supportedDocumentTypes.script]: {
+        options: ['title', 'backgroundColor'],
         dataDescription: 'The compilable JavaScript code. Use this for creating scripts.',
     },
 };
@@ -60,7 +79,7 @@ const createAnyDocumentToolParams = [
     {
         name: 'document_type',
         type: 'string',
-        description: `The type of the document to create. Supported types are: ${supportedDocumentTypes.join(', ')}`,
+        description: `The type of the document to create. Supported types are: ${Object.values(supportedDocumentTypes).join(', ')}`,
         required: true,
     },
     {
@@ -72,14 +91,11 @@ const createAnyDocumentToolParams = [
     {
         name: 'options',
         type: 'string',
-        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,
+        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;
 
@@ -87,76 +103,56 @@ type CreateAnyDocumentToolParamsType = typeof createAnyDocumentToolParams;
 
 const createAnyDocToolInfo: ToolInfo<CreateAnyDocumentToolParamsType> = {
     name: 'createAnyDocument',
-    description: `Creates any type of document (in Dash) 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
+    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(
-                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>
-        `
+                ([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>`,
+            .join('\n') +
+        `</supported_document_types>`,
     parameterRules: createAnyDocumentToolParams,
     citationRules: 'No citation needed.',
 };
 
 export class CreateAnyDocumentTool extends BaseTool<CreateAnyDocumentToolParamsType> {
-    private _addLinkedDoc: (doc_type: string, data: string | undefined, options: DocumentOptions, id: string) => void;
+    private _addLinkedDoc: (doc: parsedDoc) => Doc | undefined;
 
-    constructor(addLinkedDoc: (doc_type: string, data: string | undefined, options: DocumentOptions, id: string) => void) {
+    constructor(addLinkedDoc: (doc: parsedDoc) => Doc | undefined) {
         super(createAnyDocToolInfo);
         this._addLinkedDoc = addLinkedDoc;
     }
 
     async execute(args: ParametersType<CreateAnyDocumentToolParamsType>): Promise<Observation[]> {
         try {
-            const documentType: supportedDocumentTypesType = args.document_type.toLowerCase() as supportedDocumentTypesType;
-            let options: DocumentOptions = {};
+            const documentType = toLower(args.document_type) as unknown as supportedDocumentTypes;
+            const info = documentTypesInfo[documentType];
 
-            if (!supportedDocumentTypes.includes(documentType)) {
-                throw new Error(`Unsupported document type: ${documentType}. Supported types are: ${supportedDocumentTypes.join(', ')}.`);
+            if (info === undefined) {
+                throw new Error(`Unsupported document type: ${documentType}. Supported types are: ${Object.values(supportedDocumentTypes).join(', ')}.`);
             }
 
             if (!args.data) {
-                throw new Error(`Data is required for ${documentType} documents. ${documentTypesInfo[documentType].dataDescription}`);
+                throw new Error(`Data is required for ${documentType} documents. ${info.dataDescription}`);
             }
 
-            if (args.options) {
-                try {
-                    options = JSON.parse(args.options as string) as DocumentOptions;
-                } catch (e) {
-                    throw new Error('Options must be a valid JSON string.');
-                }
-            }
-
-            const data = args.data as string;
-            const id = uuidv4();
-
-            // Set default options if not provided
-            options.title = options.title || `New ${documentType.charAt(0).toUpperCase() + documentType.slice(1)} Document`;
+            const options: DocumentOptions = !args.options ? {} : JSON.parse(args.options);
 
-            // Call the function to add the linked document
-            this._addLinkedDoc(documentType, data, options, id);
+            // 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 ${id}.`,
-                },
-            ];
+            return [{ type: 'text', text: `Created ${documentType} document with ID ${doc?.[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 }];
         }
     }
 }