creating sorting and tagging tools for agent

1 files changed, 65 insertions, 0 deletions

diff --git a/src/client/views/nodes/chatbot/tools/SortDocsTool.ts b/src/client/views/nodes/chatbot/tools/SortDocsTool.ts
new file mode 100644
index 000000000..e91a27b36
--- /dev/null
+++ b/src/client/views/nodes/chatbot/tools/SortDocsTool.ts
@@ -0,0 +1,65 @@
+import { BaseTool } from './BaseTool';
+import { Observation } from '../types/types';
+import { ParametersType, ToolInfo } from '../types/tool_types';
+import { AgentDocumentManager } from '../utils/AgentDocumentManager';
+import { gptAPICall, GPTCallType } from '../../../../apis/gpt/GPT';
+import { ChatSortField } from '../../../collections/CollectionSubView';
+import { v4 as uuidv4 } from 'uuid';
+
+const parameterRules = [
+    {
+        name: 'sortCriteria',
+        type: 'string',
+        description: 'Criteria provided by the user to sort the documents.',
+        required: true,
+    },
+] as const;
+
+const toolInfo: ToolInfo<typeof parameterRules> = {
+    name: 'sortDocs',
+    description:
+        'Sorts documents within the current Dash environment based on user-specified criteria. Provide clear sorting criteria, such as by date, title, relevance, or custom metadata fields.',
+    parameterRules,
+    citationRules: 'No citation needed for sorting operations.',
+};
+
+export class SortDocsTool extends BaseTool<typeof parameterRules> {
+    private _docManager: AgentDocumentManager;
+
+    constructor(docManager: AgentDocumentManager) {
+        super(toolInfo);
+        this._docManager = docManager;
+        this._docManager.initializeFindDocsFreeform();
+    }
+
+    async execute(args: ParametersType<typeof parameterRules>): Promise<Observation[]> {
+        const chunkId = uuidv4();
+        try {
+            const docs = this._docManager.docIds.map((id) => this._docManager.extractDocumentMetadata(id));
+
+            const descriptions = docs
+            .filter((doc): doc is NonNullable<typeof doc> => doc !== null)
+            .map(
+                (doc) => `${doc.id}: ${doc.title} - ${doc.fields.layout.summary || ''}`
+            )
+            .join('\n');
+
+            const sortedIdsResponse = await gptAPICall(args.sortCriteria, GPTCallType.SORTDOCS, descriptions);
+            const sortedIds = sortedIdsResponse.trim().split('\n');
+
+            sortedIds.forEach((id, index) => {
+                this._docManager.editDocumentField(id, ChatSortField, index);
+            });
+
+            return [{
+                type: 'text',
+                text: `<chunk chunk_id="${chunkId}" chunk_type="sort_status">Successfully sorted ${sortedIds.length} documents based on "${args.sortCriteria}".</chunk>`,
+            }];
+        } catch (error) {
+            return [{
+                type: 'text',
+                text: `<chunk chunk_id="${chunkId}" chunk_type="error">Sorting failed: ${error instanceof Error ? error.message : String(error)}</chunk>`,
+            }];
+        }
+    }
+}