fixes for chatbox that already exists when a collection is opened to still have access to the collection's docs.

1 files changed, 268 insertions, 300 deletions

diff --git a/src/client/views/nodes/chatbot/tools/FilterDocsTool.ts b/src/client/views/nodes/chatbot/tools/FilterDocsTool.ts
index b160badde..a921f6058 100644
--- a/src/client/views/nodes/chatbot/tools/FilterDocsTool.ts
+++ b/src/client/views/nodes/chatbot/tools/FilterDocsTool.ts
@@ -3,343 +3,311 @@ import { BaseTool } from './BaseTool';
 import { Observation } from '../types/types';
 import { ParametersType, ToolInfo } from '../types/tool_types';
 import { AgentDocumentManager } from '../utils/AgentDocumentManager';
-import {
-  gptAPICall,
-  GPTCallType,
-  GPTDocCommand,
-  DescriptionSeperator,
-  DataSeperator,
-  DocSeperator,
-} from '../../../../apis/gpt/GPT';
+import { gptAPICall, GPTCallType, GPTDocCommand, DescriptionSeperator, DataSeperator, DocSeperator } from '../../../../apis/gpt/GPT';
 import { v4 as uuidv4 } from 'uuid';
 import { TagItem } from '../../../TagsView';
 import { DocumentView } from '../../DocumentView';
 import { Doc } from '../../../../../fields/Doc';
+import { computed } from 'mobx';
 
 const parameterRules = [
-  {
-    name: 'filterCriteria',
-    type: 'string',
-    description: 'Natural-language criteria for choosing a subset of documents (e.g., "show me only research papers about AI", "filter for documents created last week")',
-    required: true,
-  },
+    {
+        name: 'filterCriteria',
+        type: 'string',
+        description: 'Natural-language criteria for choosing a subset of documents (e.g., "show me only research papers about AI", "filter for documents created last week")',
+        required: true,
+    },
 ] as const;
 
 const toolInfo: ToolInfo<typeof parameterRules> = {
-  name: 'filterDocs',
-  description: 'Filters documents in a collection based on user-specified natural-language criteria. Replaces GPTPopup filtering functionality with enhanced error handling and management features.',
-  parameterRules,
-  citationRules: 'No citation needed for filtering operations.',
+    name: 'filterDocs',
+    description: 'Filters documents in a collection based on user-specified natural-language criteria. Replaces GPTPopup filtering functionality with enhanced error handling and management features.',
+    parameterRules,
+    citationRules: 'No citation needed for filtering operations.',
 };
 
 export class FilterDocsTool extends BaseTool<typeof parameterRules> {
-  private _docManager: AgentDocumentManager;
-  static ChatTag = '#chat'; // tag used to mark filtered documents
-  private _collectionView: DocumentView;
-  private _documentDescriptions: Promise<string> | undefined;
-  private _textToDocMap = new Map<string, Doc>();
-  private _lastFilterCriteria: string | undefined;
-
-  constructor(docManager: AgentDocumentManager, collectionView: DocumentView) {
-    super(toolInfo);
-    this._docManager = docManager;
-            this._docManager.initializeDocuments();
-    this._collectionView = collectionView;
-    this._initializeDocumentContext();
-  }
-
-  /**
-   * Initialize document context similar to GPTPopup's componentDidMount behavior
-   */
-  private _initializeDocumentContext() {
-    const selectedDoc = this._collectionView?.Document;
-    // Use any type to avoid complex type checking while maintaining runtime safety
-    const componentView = selectedDoc?.ComponentView as any;
-    const hasChildDocs = componentView?.hasChildDocs;
-    
-    if (hasChildDocs && typeof hasChildDocs === 'function') {
-      this._textToDocMap.clear();
-      try {
-        const childDocs = hasChildDocs() as Doc[];
-        this._documentDescriptions = Promise.all(
-          childDocs.map((doc: Doc) =>
-            Doc.getDescription(doc).then(text => {
-              const cleanText = text.replace(/\n/g, ' ').trim();
-              this._textToDocMap.set(cleanText, doc);
-              return `${DescriptionSeperator}${cleanText}${DescriptionSeperator}`;
-            })
-          )
-        ).then(docDescriptions => docDescriptions.join(''));
-      } catch (error) {
-        console.warn('[FilterDocsTool] Error initializing document context:', error);
-      }
+    private _docManager: AgentDocumentManager;
+    static ChatTag = '#chat'; // tag used to mark filtered documents
+    private _collectionView: DocumentView;
+    private _documentDescriptions: Promise<string> | undefined;
+    private _lastFilterCriteria: string | undefined;
+
+    constructor(docManager: AgentDocumentManager, collectionView: DocumentView) {
+        super(toolInfo);
+        this._docManager = docManager;
+        this._docManager.initializeDocuments();
+        this._collectionView = collectionView;
     }
-  }
-
-  /**
-   * Check if a document filter is currently active
-   */
-  static hasActiveFilter(parentDoc: Doc | undefined): boolean {
-    if (!parentDoc) return false;
-    const result = Doc.hasDocFilter(parentDoc, 'tags', FilterDocsTool.ChatTag);
-    return Boolean(result);
-  }
-
-  /**
-   * Clear all chat filters from a collection (equivalent to GPTPopup's clear filter toggle)
-   */
-  static clearChatFilter(parentDoc: Doc | undefined, childDocs?: Doc[]): boolean {
-    if (!parentDoc) return false;
-
-    try {
-      // Remove filter from parent document
-      Doc.setDocFilter(parentDoc, 'tags', FilterDocsTool.ChatTag, 'remove');
-      
-      // Remove tags from child documents
-      let docsToProcess: Doc[] = [];
-      if (childDocs) {
-        docsToProcess = childDocs;
-      } else {
-        // Try to get child docs from ComponentView
-        const componentView = parentDoc.ComponentView as any;
-        const hasChildDocs = componentView?.hasChildDocs;
-        if (hasChildDocs && typeof hasChildDocs === 'function') {
-          docsToProcess = hasChildDocs() as Doc[];
+
+    @computed 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 `${DescriptionSeperator}${cleanText}${DescriptionSeperator}`;
+                        })
+                    )
+                ).then(docDescriptions => docDescriptions.join(''));
+                return textToDocMap;
+            } catch (error) {
+                console.warn('[FilterDocsTool] Error initializing document context:', error);
+            }
         }
-      }
-      
-      docsToProcess.forEach((doc: Doc) => TagItem.removeTagFromDoc(doc, FilterDocsTool.ChatTag));
-      
-      return true;
-    } catch (error) {
-      console.error('[FilterDocsTool] Error clearing chat filter:', error);
-      return false;
+        return undefined;
     }
-  }
-
-  /**
-   * Determine if user input is a filter command using GPT
-   */
-  async isFilterCommand(userPrompt: string): Promise<boolean> {
-    try {
-      const commandType = await gptAPICall(userPrompt, GPTCallType.COMMANDTYPE, undefined, true);
-      const commandNumber = parseInt(commandType.split(' ')[0][0]);
-      return commandNumber === GPTDocCommand.Filter;
-    } catch (error) {
-      console.error('[FilterDocsTool] Error determining command type:', error);
-      return false;
+
+    /**
+     * Check if a document filter is currently active
+     */
+    static hasActiveFilter(parentDoc: Doc | undefined): boolean {
+        if (!parentDoc) return false;
+        const result = Doc.hasDocFilter(parentDoc, 'tags', FilterDocsTool.ChatTag);
+        return Boolean(result);
     }
-  }
-
-  async execute(args: ParametersType<typeof parameterRules>): Promise<Observation[]> {
-    const chunkId = uuidv4();
-    this._lastFilterCriteria = args.filterCriteria;
-    console.log('[FilterDocsTool] Executing filter with criteria:', args.filterCriteria);
-
-    // Validate parent view and document
-    const parentView = this._collectionView;
-    const parentDoc = parentView?.Document;
-    if (!parentDoc) {
-      console.error('[FilterDocsTool] Missing parent DocumentView/Document!');
-      return [
-        {
-          type: 'text',
-          text: `<chunk chunk_id="${chunkId}" chunk_type="error">
-FilterDocsTool: No parent collection document found. Please ensure you're working within a document collection.
-</chunk>`,
-        },
-      ];
+
+    /**
+     * Clear all chat filters from a collection (equivalent to GPTPopup's clear filter toggle)
+     */
+    static clearChatFilter(parentDoc: Doc | undefined, childDocs?: Doc[]): boolean {
+        if (!parentDoc) return false;
+
+        try {
+            // Remove filter from parent document
+            Doc.setDocFilter(parentDoc, 'tags', FilterDocsTool.ChatTag, 'remove');
+
+            // Remove tags from child documents
+            let docsToProcess: Doc[] = [];
+            if (childDocs) {
+                docsToProcess = childDocs;
+            } else {
+                // Try to get child docs from ComponentView
+                const componentView = parentDoc.ComponentView as any;
+                const hasChildDocs = componentView?.hasChildDocs;
+                if (hasChildDocs && typeof hasChildDocs === 'function') {
+                    docsToProcess = hasChildDocs() as Doc[];
+                }
+            }
+
+            docsToProcess.forEach((doc: Doc) => TagItem.removeTagFromDoc(doc, FilterDocsTool.ChatTag));
+
+            return true;
+        } catch (error) {
+            console.error('[FilterDocsTool] Error clearing chat filter:', error);
+            return false;
+        }
     }
 
-    try {
-      // Method 1: Use pre-computed document descriptions if available (from componentDidMount-like behavior)
-      let prompt: string;
-      let textToDocMap: Map<string, Doc>;
-
-      if (this._documentDescriptions && this._textToDocMap.size > 0) {
-        console.log('[FilterDocsTool] Using pre-computed document descriptions');
-        prompt = await this._documentDescriptions;
-        textToDocMap = this._textToDocMap;
-      } else {
-        // Method 2: Build descriptions from scratch using docManager
-        console.log('[FilterDocsTool] Building document descriptions from docManager');
-        textToDocMap = new Map<string, Doc>();
-        const blocks: string[] = [];
-
-        for (const id of this._docManager.docIds) {
-          const descRaw = await this._docManager.getDocDescription(id);
-          const desc = descRaw.replace(/\n/g, ' ').trim();
-
-          if (!desc) {
-            console.warn(`[FilterDocsTool] Skipping document ${id} with empty description`);
-            continue;
-          }
-
-          const doc = this._docManager.getDocument(id);
-          if (doc) {
-            textToDocMap.set(desc, doc);
-            blocks.push(`${DescriptionSeperator}${desc}${DescriptionSeperator}`);
-          }
+    /**
+     * Determine if user input is a filter command using GPT
+     */
+    async isFilterCommand(userPrompt: string): Promise<boolean> {
+        try {
+            const commandType = await gptAPICall(userPrompt, GPTCallType.COMMANDTYPE, undefined, true);
+            const commandNumber = parseInt(commandType.split(' ')[0][0]);
+            return commandNumber === GPTDocCommand.Filter;
+        } catch (error) {
+            console.error('[FilterDocsTool] Error determining command type:', error);
+            return false;
         }
+    }
 
-        prompt = blocks.join('');
-      }
+    async execute(args: ParametersType<typeof parameterRules>): Promise<Observation[]> {
+        const chunkId = uuidv4();
+        this._lastFilterCriteria = args.filterCriteria;
+        console.log('[FilterDocsTool] Executing filter with criteria:', args.filterCriteria);
+
+        // Validate parent view and document
+        const parentView = this._collectionView;
+        const parentDoc = parentView?.Document;
+        if (!parentDoc) {
+            console.error('[FilterDocsTool] Missing parent DocumentView/Document!');
+            return [
+                {
+                    type: 'text',
+                    text: `<chunk chunk_id="${chunkId}" chunk_type="error">
+FilterDocsTool: No parent collection document found. Please ensure you're working within a document collection.
+</chunk>`,
+                },
+            ];
+        }
 
-      if (!prompt || textToDocMap.size === 0) {
-        return [
-          {
-            type: 'text',
-            text: `<chunk chunk_id="${chunkId}" chunk_type="error">
+        try {
+            // Method 1: Use pre-computed document descriptions if available (from componentDidMount-like behavior)
+            let prompt: string;
+            let textToDocMap = await this.TextToDocMap;
+
+            if (textToDocMap && textToDocMap.size > 0 && this._documentDescriptions) {
+                console.log('[FilterDocsTool] Using pre-computed document descriptions');
+                prompt = await this._documentDescriptions;
+            } else {
+                // Method 2: Build descriptions from scratch using docManager
+                console.log('[FilterDocsTool] Building document descriptions from docManager');
+                textToDocMap = new Map<string, Doc>();
+                const blocks: string[] = [];
+
+                for (const id of this._docManager.docIds) {
+                    const descRaw = await this._docManager.getDocDescription(id);
+                    const desc = descRaw.replace(/\n/g, ' ').trim();
+
+                    if (!desc) {
+                        console.warn(`[FilterDocsTool] Skipping document ${id} with empty description`);
+                        continue;
+                    }
+
+                    const doc = this._docManager.getDocument(id);
+                    if (doc) {
+                        textToDocMap.set(desc, doc);
+                        blocks.push(`${DescriptionSeperator}${desc}${DescriptionSeperator}`);
+                    }
+                }
+
+                prompt = blocks.join('');
+            }
+
+            if (!prompt || textToDocMap?.size === 0) {
+                return [
+                    {
+                        type: 'text',
+                        text: `<chunk chunk_id="${chunkId}" chunk_type="error">
 No documents found to filter. The collection appears to be empty or documents lack descriptions.
 </chunk>`,
-          },
-        ];
-      }
-
-      console.log(`[FilterDocsTool] Processing ${textToDocMap.size} documents with filter criteria`);
-
-      // Call GPT for document subset selection
-      const gptResponse = await gptAPICall(args.filterCriteria, GPTCallType.SUBSETDOCS, prompt);
-      console.log('[FilterDocsTool] GPT response received:', gptResponse.substring(0, 200) + '...');
-
-      // Clear existing filters using the enhanced clear method
-      const childDocs = Array.from(textToDocMap.values());
-      FilterDocsTool.clearChatFilter(parentDoc, childDocs);
-
-      // Parse GPT response and apply new filters
-      const filteredResults = this._parseGptResponseAndApplyFilters(
-        gptResponse,
-        textToDocMap,
-        parentDoc,
-        chunkId
-      );
-
-      return filteredResults;
-
-    } catch (error) {
-      console.error('[FilterDocsTool] Execution error:', error);
-      return [
-        {
-          type: 'text',
-          text: `<chunk chunk_id="${chunkId}" chunk_type="error">
+                    },
+                ];
+            }
+
+            console.log(`[FilterDocsTool] Processing ${textToDocMap?.size} documents with filter criteria`);
+
+            // Call GPT for document subset selection
+            const gptResponse = await gptAPICall(args.filterCriteria, GPTCallType.SUBSETDOCS, prompt);
+            console.log('[FilterDocsTool] GPT response received:', gptResponse.substring(0, 200) + '...');
+
+            // Clear existing filters using the enhanced clear method
+            const childDocs = Array.from(textToDocMap?.values() ?? []);
+            FilterDocsTool.clearChatFilter(parentDoc, childDocs);
+
+            // Parse GPT response and apply new filters
+            const filteredResults = this._parseGptResponseAndApplyFilters(gptResponse, textToDocMap, parentDoc, chunkId);
+
+            return filteredResults;
+        } catch (error) {
+            console.error('[FilterDocsTool] Execution error:', error);
+            return [
+                {
+                    type: 'text',
+                    text: `<chunk chunk_id="${chunkId}" chunk_type="error">
 Filtering failed: ${error instanceof Error ? error.message : String(error)}
 Please try rephrasing your filter criteria or check if the collection contains valid documents.
 </chunk>`,
-        },
-      ];
-    }
-  }
-
-  /**
-   * Parse GPT response and apply filters with enhanced error handling
-   */
-  private _parseGptResponseAndApplyFilters(
-    gptResponse: string,
-    textToDocMap: Map<string, Doc>,
-    parentDoc: Doc,
-    chunkId: string
-  ): Observation[] {
-    const filteredDocIds: string[] = [];
-    const errors: string[] = [];
-    let explanation = '';
-
-    try {
-      // Extract explanation if present (surrounded by DocSeperator)
-      const explanationMatch = gptResponse.match(new RegExp(`${DocSeperator}\\s*([\\s\\S]*?)\\s*(?:${DocSeperator}|$)`));
-      explanation = explanationMatch?.[1]?.trim() || '';
-
-      // Parse document descriptions from GPT response
-      const docBlocks = gptResponse
-        .split(DescriptionSeperator)
-        .filter(block => block.trim() !== '')
-        .map(block => block.replace(/\n/g, ' ').trim());
-
-      console.log(`[FilterDocsTool] Found ${docBlocks.length} document blocks in GPT response`);
-
-      docBlocks.forEach((block, index) => {
-        // Split by DataSeperator if present (for future extensibility)
-        const [descText = ''] = block.split(DataSeperator).map(s => s.trim());
-        
-        if (!descText) {
-          console.warn(`[FilterDocsTool] Skipping empty description block ${index}`);
-          return;
-        }
-
-        const doc = textToDocMap.get(descText);
-        if (!doc) {
-          console.warn(`[FilterDocsTool] No matching document found for: "${descText.substring(0, 50)}..."`);
-          errors.push(`No match found for: "${descText.substring(0, 30)}..."`);
-          return;
+                },
+            ];
         }
+    }
 
-        // Apply the chat tag to mark as filtered
-        try {
-          TagItem.addTagToDoc(doc, FilterDocsTool.ChatTag);
-          const docId = String(doc.id || `doc_${index}`);
-          filteredDocIds.push(docId);
-          console.log(`[FilterDocsTool] Tagged document: ${docId}`);
-        } catch (tagError) {
-          console.error(`[FilterDocsTool] Error tagging document:`, tagError);
-          errors.push(`Failed to tag document: ${doc.id || 'unknown'}`);
-        }
-      });
+    /**
+     * Parse GPT response and apply filters with enhanced error handling
+     */
+    private _parseGptResponseAndApplyFilters(gptResponse: string, textToDocMap: Map<string, Doc>, parentDoc: Doc, chunkId: string): Observation[] {
+        const filteredDocIds: string[] = [];
+        const errors: string[] = [];
+        let explanation = '';
 
-      // Apply the document filter to show only tagged documents
-      if (filteredDocIds.length > 0) {
         try {
-          Doc.setDocFilter(parentDoc, 'tags', FilterDocsTool.ChatTag, 'check');
-          const parentDocId = String(parentDoc.id || 'unknown');
-          console.log(`[FilterDocsTool] Applied filter to parent document: ${parentDocId}`);
-        } catch (filterError) {
-          console.error('[FilterDocsTool] Error setting document filter:', filterError);
-          errors.push('Failed to apply collection filter');
+            // Extract explanation if present (surrounded by DocSeperator)
+            const explanationMatch = gptResponse.match(new RegExp(`${DocSeperator}\\s*([\\s\\S]*?)\\s*(?:${DocSeperator}|$)`));
+            explanation = explanationMatch?.[1]?.trim() || '';
+
+            // Parse document descriptions from GPT response
+            const docBlocks = gptResponse
+                .split(DescriptionSeperator)
+                .filter(block => block.trim() !== '')
+                .map(block => block.replace(/\n/g, ' ').trim());
+
+            console.log(`[FilterDocsTool] Found ${docBlocks.length} document blocks in GPT response`);
+
+            docBlocks.forEach((block, index) => {
+                // Split by DataSeperator if present (for future extensibility)
+                const [descText = ''] = block.split(DataSeperator).map(s => s.trim());
+
+                if (!descText) {
+                    console.warn(`[FilterDocsTool] Skipping empty description block ${index}`);
+                    return;
+                }
+
+                const doc = textToDocMap.get(descText);
+                if (!doc) {
+                    console.warn(`[FilterDocsTool] No matching document found for: "${descText.substring(0, 50)}..."`);
+                    errors.push(`No match found for: "${descText.substring(0, 30)}..."`);
+                    return;
+                }
+
+                // Apply the chat tag to mark as filtered
+                try {
+                    TagItem.addTagToDoc(doc, FilterDocsTool.ChatTag);
+                    const docId = String(doc.id || `doc_${index}`);
+                    filteredDocIds.push(docId);
+                    console.log(`[FilterDocsTool] Tagged document: ${docId}`);
+                } catch (tagError) {
+                    console.error(`[FilterDocsTool] Error tagging document:`, tagError);
+                    errors.push(`Failed to tag document: ${doc.id || 'unknown'}`);
+                }
+            });
+
+            // Apply the document filter to show only tagged documents
+            if (filteredDocIds.length > 0) {
+                try {
+                    Doc.setDocFilter(parentDoc, 'tags', FilterDocsTool.ChatTag, 'check');
+                    const parentDocId = String(parentDoc.id || 'unknown');
+                    console.log(`[FilterDocsTool] Applied filter to parent document: ${parentDocId}`);
+                } catch (filterError) {
+                    console.error('[FilterDocsTool] Error setting document filter:', filterError);
+                    errors.push('Failed to apply collection filter');
+                }
+            }
+        } catch (parseError) {
+            console.error('[FilterDocsTool] Error parsing GPT response:', parseError);
+            errors.push(`Response parsing failed: ${parseError instanceof Error ? parseError.message : String(parseError)}`);
         }
-      }
 
-    } catch (parseError) {
-      console.error('[FilterDocsTool] Error parsing GPT response:', parseError);
-      errors.push(`Response parsing failed: ${parseError instanceof Error ? parseError.message : String(parseError)}`);
-    }
-
-    // Build result message  
-    const filterCriteria = this._lastFilterCriteria || 'the specified criteria';
-    const successMessage = filteredDocIds.length > 0 
-      ? `Successfully filtered ${filteredDocIds.length} document(s) matching "${filterCriteria}".`
-      : 'No documents matched the specified filter criteria.';
+        // Build result message
+        const filterCriteria = this._lastFilterCriteria || 'the specified criteria';
+        const successMessage = filteredDocIds.length > 0 ? `Successfully filtered ${filteredDocIds.length} document(s) matching "${filterCriteria}".` : 'No documents matched the specified filter criteria.';
 
-    const explanationMessage = explanation 
-      ? `\n\nGPT Explanation: ${explanation}`
-      : '';
+        const explanationMessage = explanation ? `\n\nGPT Explanation: ${explanation}` : '';
 
-    const errorMessage = errors.length > 0 
-      ? `\n\nWarnings: ${errors.join('; ')}`
-      : '';
+        const errorMessage = errors.length > 0 ? `\n\nWarnings: ${errors.join('; ')}` : '';
 
-    return [
-      {
-        type: 'text',
-        text: `<chunk chunk_id="${chunkId}" chunk_type="filter_result">
+        return [
+            {
+                type: 'text',
+                text: `<chunk chunk_id="${chunkId}" chunk_type="filter_result">
 ${successMessage}${explanationMessage}${errorMessage}
 
 Filtered document IDs: ${filteredDocIds.length > 0 ? filteredDocIds.join(', ') : '(none)'}
 Filter tag applied: ${FilterDocsTool.ChatTag}
 </chunk>`,
-      },
-    ];
-  }
-
-  /**
-   * Get current filter status for UI integration
-   */
-  getFilterStatus(): { isActive: boolean; documentCount: number; filterTag: string } {
-    const parentDoc = this._collectionView?.Document;
-    const isActive = FilterDocsTool.hasActiveFilter(parentDoc);
-    const documentCount = this._textToDocMap.size;
-    
-    return {
-      isActive,
-      documentCount,
-      filterTag: FilterDocsTool.ChatTag,
-    };
-  }
-} 
\ No newline at end of file
+            },
+        ];
+    }
+
+    /**
+     * Get current filter status for UI integration
+     */
+    getFilterStatus(): { isActive: boolean; documentCount: number; filterTag: string } {
+        const parentDoc = this._collectionView?.Document;
+        const isActive = FilterDocsTool.hasActiveFilter(parentDoc);
+        const documentCount = this.TextToDocMap?.size ?? 0;
+
+        return {
+            isActive,
+            documentCount,
+            filterTag: FilterDocsTool.ChatTag,
+        };
+    }
+}