1 files changed, 108 insertions, 67 deletions

diff --git a/src/client/views/nodes/chatbot/tools/SortDocsTool.ts b/src/client/views/nodes/chatbot/tools/SortDocsTool.ts
index 9eb0cf4d1..a1f43a8f2 100644
--- a/src/client/views/nodes/chatbot/tools/SortDocsTool.ts
+++ b/src/client/views/nodes/chatbot/tools/SortDocsTool.ts
@@ -2,97 +2,138 @@ import { BaseTool } from './BaseTool';
 import { Observation } from '../types/types';
 import { ParametersType, ToolInfo } from '../types/tool_types';
 import { AgentDocumentManager } from '../utils/AgentDocumentManager';
-import { gptAPICall, GPTCallType, DescriptionSeperator } from '../../../../apis/gpt/GPT';
+import { gptAPICall, GPTCallType, DescStart, DescEnd } from '../../../../apis/gpt/GPT';
 import { ChatSortField } from '../../../collections/CollectionSubView';
 import { v4 as uuidv4 } from 'uuid';
 import { DocumentView } from '../../DocumentView';
 import { docSortings } from '../../../collections/CollectionSubView';
-
+import { Doc } from '../../../../../fields/Doc';
 
 const parameterRules = [
-  {
-    name: 'sortCriteria',
-    type: 'string',
-    description: 'Criteria provided by the user to sort the documents.',
-    required: true,
-  },
+    {
+        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.',
-  parameterRules,
-  citationRules: 'No citation needed for sorting operations.',
+    name: 'sortDocs',
+    description: 'Sorts documents within the current Dash environment based on user-specified criteria.',
+    parameterRules,
+    citationRules: 'No citation needed for sorting operations.',
 };
 
 export class SortDocsTool extends BaseTool<typeof parameterRules> {
-  private _docManager: AgentDocumentManager;
-  private _collectionView: DocumentView;
+    private _docManager: AgentDocumentManager;
+    private _collectionView: DocumentView;
+    private _documentDescriptions: Promise<string> | undefined;
+
+    constructor(docManager: AgentDocumentManager, collectionView: DocumentView) {
+        super(toolInfo);
+        // Grab the parent collection’s DocumentView (the ChatBox container)
+        // We assume the ChatBox itself is currently selected in its parent view.
+        this._collectionView = collectionView;
+        this._docManager = docManager;
+        this._docManager.initializeDocuments();
+    }
+
+    get TextToDocMap() {
+        // Use any type to avoid complex type checking while maintaining runtime safety
+        const childDocs = this._collectionView?.ComponentView?.hasChildDocs?.();
+        if (childDocs) {
+            const textToDocMap = new Map<string, Doc>();
+            try {
+                this._documentDescriptions = Promise.all(
+                    childDocs.map((doc: Doc) =>
+                        Doc.getDescription(doc).then(text => {
+                            const cleanText = text.replace(/\n/g, ' ').trim();
+                            textToDocMap.set(cleanText, doc);
+                            return `${DescStart}${cleanText}${DescEnd}`;
+                        })
+                    )
+                ).then(docDescriptions => docDescriptions.join(''));
+                return textToDocMap;
+            } catch (error) {
+                console.warn('[SortDocsTool] Error initializing document context:', error);
+            }
+        }
+        return undefined;
+    }
+
+    async execute(args: ParametersType<typeof parameterRules>): Promise<Observation[]> {
+        const chunkId = uuidv4();
+        let textToDocMap = await this.TextToDocMap;
+        await this._documentDescriptions;
+        let chunks: string;
 
-  constructor(docManager: AgentDocumentManager, collectionView: DocumentView)
-  {
-    super(toolInfo);
-     // Grab the parent collection’s DocumentView (the ChatBox container)
-    // We assume the ChatBox itself is currently selected in its parent view.
-    this._collectionView = collectionView;
-    this._docManager = docManager;
-            this._docManager.initializeDocuments();
-  }
+        if (textToDocMap && textToDocMap.size > 0 && this._documentDescriptions) {
+            console.log('[SortDocsTool] Using pre-computed document descriptions');
+            chunks = await this._documentDescriptions;
+        } else {
+            // Method 2: Build descriptions from scratch using docManager
+            console.log('[SortDocsTool] Building document descriptions from docManager');
+            textToDocMap = new Map<string, Doc>();
+            const blocks: string[] = [];
 
-  async execute(args: ParametersType<typeof parameterRules>): Promise<Observation[]> {
-    const chunkId = uuidv4();
+            for (const id of this._docManager.docIds) {
+                const descRaw = await this._docManager.getDocDescription(id);
+                const desc = descRaw.replace(/\n/g, ' ').trim();
 
-    // 1) gather metadata & build map from text→id
- const textToId = new Map<string, string>();
+                if (!desc) {
+                    console.warn(`[SortDocsTool] Skipping document ${id} with empty description`);
+                    continue;
+                }
 
-    const chunks = (await Promise.all(
-    this._docManager.docIds.map(async id => {
-        const text = await this._docManager.getDocDescription(id);
-        textToId.set(text,id);
-        return DescriptionSeperator + text + DescriptionSeperator;
-    })
-    ))
-    .join('');
-    try {
-      // 2) call GPT to sort those chunks
-      const gptResponse = await gptAPICall(args.sortCriteria, GPTCallType.SORTDOCS, chunks);
-      console.log('GPT RESP:', gptResponse);
+                const doc = this._docManager.getDocument(id);
+                if (doc) {
+                    textToDocMap.set(desc, doc);
+                    blocks.push(`${DescStart}${desc}${DescEnd}`);
+                }
+            }
 
-      // 3) parse & map back to IDs
-      const sortedIds = gptResponse
-        .split(DescriptionSeperator)
-        .filter(s => s.trim() !== '')
-        .map(s => s.replace(/\n/g, ' ').trim())
-        .map(s => textToId.get(s))         // lookup in our map
-        .filter((id): id is string => !!id);
+            chunks = blocks.join('');
+        }
+        try {
+            // 2) call GPT to sort those chunks
+            const gptResponse = await gptAPICall(args.sortCriteria, GPTCallType.SORTDOCS, chunks);
+            console.log('GPT RESP:', gptResponse);
 
-      // 4) write back the ordering
-      sortedIds.forEach((docId, idx) => {
-        this._docManager.editDocumentField(docId, ChatSortField, idx);
-      });
+            // 3) parse & map back to IDs
+            const sortedIds = gptResponse
+                .split(DescEnd)
+                .filter(s => s.trim() !== '')
+                .map(s => s.replace(DescStart, '').replace(/\n/g, ' ').trim())
+                .map(s => textToDocMap.get(s)) // lookup in our map
+                .filter(doc => doc)
+                .map(doc => doc!);
 
-       const fieldKey = this._collectionView.ComponentView!.fieldKey;
-         this._collectionView.Document[ `${fieldKey}_sort` ] = docSortings.Chat;
+            // 4) write back the ordering
+            sortedIds.forEach((doc, idx) => {
+                doc[ChatSortField] = idx;
+            });
 
+            const fieldKey = this._collectionView.ComponentView!.fieldKey;
+            this._collectionView.Document[`${fieldKey}_sort`] = docSortings.Chat;
 
-      return [
-        {
-          type: 'text',
-          text: `<chunk chunk_id="${chunkId}" chunk_type="sort_status">
+            return [
+                {
+                    type: 'text',
+                    text: `<chunk chunk_id="${chunkId}" chunk_type="sort_status">
 Successfully sorted ${sortedIds.length} documents by "${args.sortCriteria}".
 </chunk>`,
-        },
-      ];
-    } catch (err) {
-      return [
-        {
-          type: 'text',
-          text: `<chunk chunk_id="${chunkId}" chunk_type="error">
+                },
+            ];
+        } catch (err) {
+            return [
+                {
+                    type: 'text',
+                    text: `<chunk chunk_id="${chunkId}" chunk_type="error">
 Sorting failed: ${err instanceof Error ? err.message : err}
 </chunk>`,
-        },
-      ];
+                },
+            ];
+        }
     }
-  }
 }