1 files changed, 23 insertions, 15 deletions

diff --git a/src/client/views/nodes/chatbot/tools/GetDocsTool.ts b/src/client/views/nodes/chatbot/tools/GetDocsTool.ts
index b6e549d93..26756522c 100644
--- a/src/client/views/nodes/chatbot/tools/GetDocsTool.ts
+++ b/src/client/views/nodes/chatbot/tools/GetDocsTool.ts
@@ -1,37 +1,45 @@
 import { Observation } from '../types/types';
+import { ParametersType } from './ToolTypes';
 import { BaseTool } from './BaseTool';
 import { DocServer } from '../../../../DocServer';
 import { Docs } from '../../../../documents/Documents';
 import { DocumentView } from '../../DocumentView';
 import { OpenWhere } from '../../OpenWhere';
+import { DocCast } from '../../../../../fields/Types';
 
-export class GetDocsTool extends BaseTool<{ title: { type: string; description: string; required: boolean }; document_ids: { type: string; description: string; required: boolean } }> {
+const getDocsToolParams = [
+    {
+        name: 'title',
+        type: 'string',
+        description: 'Title of the collection being created from retrieved documents',
+        required: true,
+    },
+    {
+        name: 'document_ids',
+        type: 'string[]',
+        description: 'List of document IDs to retrieve',
+        required: true,
+    },
+] as const;
+
+type GetDocsToolParamsType = typeof getDocsToolParams;
+
+export class GetDocsTool extends BaseTool<GetDocsToolParamsType> {
     private _docView: DocumentView;
 
     constructor(docView: DocumentView) {
         super(
             'retrieveDocs',
             'Retrieves the contents of all Documents that the user is interacting with in Dash',
-            {
-                title: {
-                    type: 'string',
-                    description: 'Title of the collection being created from retrieved documents',
-                    required: true,
-                },
-                document_ids: {
-                    type: 'string[]',
-                    description: 'List of document IDs to retrieve',
-                    required: true,
-                },
-            },
+            getDocsToolParams,
             'No need to provide anything. Just run the tool and it will retrieve the contents of all Documents that the user is interacting with in Dash.',
             'Returns the documents in Dash in JSON form.'
         );
         this._docView = docView;
     }
 
-    async execute(args: { title: string; document_ids: string[] }): Promise<Observation[]> {
-        const docs = args.document_ids.map(doc_id => DocServer.GetCachedRefField(doc_id));
+    async execute(args: ParametersType<GetDocsToolParamsType>): Promise<Observation[]> {
+        const docs = args.document_ids.map(doc_id => DocCast(DocServer.GetCachedRefField(doc_id)));
         const collection = Docs.Create.FreeformDocument(docs, { title: args.title });
         this._docView._props.addDocTab(collection, OpenWhere.addRight);
         return [{ type: 'text', text: `Collection created in Dash called ${args.title}` }];