Trilium/apps/server/src/services/llm/tools/search_notes_tool.ts

/**
 * Search Notes Tool
 *
 * This tool allows the LLM to search for notes using semantic search.
 */

import type { Tool, ToolHandler } from './tool_interfaces.js';
import log from '../../log.js';
import aiServiceManager from '../ai_service_manager.js';
import becca from '../../../becca/becca.js';
import { ContextExtractor } from '../context/index.js';

/**
 * Definition of the search notes tool
 */
export const searchNotesToolDefinition: Tool = {
    type: 'function',
    function: {
        name: 'search_notes',
        description: `SEMANTIC/CONCEPTUAL search for notes. Finds notes related to concepts, topics, or themes even without exact keyword matches.
        
        BEST FOR: Finding notes about ideas, concepts, or topics described in various ways
        USE WHEN: Looking for conceptual relationships, thematic content, or related ideas
        DIFFERENT FROM: keyword_search (which finds exact text matches)
        
        TIPS: 
        - Use descriptive phrases like "project management methodologies" rather than single words
        - Think conceptually: "machine learning classification" vs just "ML"
        - Results include noteId values - ALWAYS use these IDs (not titles) with other tools
        
        NEXT STEPS: Use read_note with returned noteId values to get full content`,
        parameters: {
            type: 'object',
            properties: {
                query: {
                    type: 'string',
                    description: `Descriptive search query for semantic matching.
                    
                    GOOD EXAMPLES:
                    - "machine learning algorithms for classification"
                    - "personal productivity and time management techniques"
                    - "software development best practices"
                    
                    AVOID:
                    - Single words: "ML", "productivity"
                    - Overly broad: "work", "notes"
                    - Overly specific: exact phrases that might not exist`
                },
                parentNoteId: {
                    type: 'string',
                    description: `SCOPE LIMITER: Search only within children of this note.
                    
                    IMPORTANT: Must be a noteId (like "abc123def456") from previous search results - NOT a note title.
                    
                    USE FOR: Searching within specific projects, categories, or sections.`
                },
                maxResults: {
                    type: 'number',
                    description: 'Number of results (1-20, default: 5). Use 10-15 for comprehensive exploration, 3-5 for quick lookup.'
                },
                summarize: {
                    type: 'boolean',
                    description: 'AI SUMMARIES: Get intelligent summaries instead of truncated text (default: false). Use true for cleaner result overview.'
                }
            },
            required: ['query']
        }
    }
};

/**
 * Get or create the vector search tool dependency
 * @returns The vector search tool or null if it couldn't be created
 */
async function getOrCreateVectorSearchTool(): Promise<any> {
    try {
        // Try to get the existing vector search tool
        let vectorSearchTool = aiServiceManager.getVectorSearchTool();

        if (vectorSearchTool) {
            log.info(`Found existing vectorSearchTool`);
            return vectorSearchTool;
        }

        // No existing tool, try to initialize it
        log.info(`VectorSearchTool not found, attempting initialization`);

        // Get agent tools manager and initialize it
        const agentTools = aiServiceManager.getAgentTools();
        if (agentTools && typeof agentTools.initialize === 'function') {
            try {
                // Force initialization to ensure it runs even if previously marked as initialized
                await agentTools.initialize(true);
            } catch (initError: any) {
                log.error(`Failed to initialize agent tools: ${initError.message}`);
                return null;
            }
        } else {
            log.error('Agent tools manager not available');
            return null;
        }

        // Try getting the vector search tool again after initialization
        vectorSearchTool = aiServiceManager.getVectorSearchTool();

        if (vectorSearchTool) {
            log.info('Successfully created vectorSearchTool');
            return vectorSearchTool;
        } else {
            log.error('Failed to create vectorSearchTool after initialization');
            return null;
        }
    } catch (error: any) {
        log.error(`Error getting or creating vectorSearchTool: ${error.message}`);
        return null;
    }
}

/**
 * Search notes tool implementation
 */
export class SearchNotesTool implements ToolHandler {
    public definition: Tool = searchNotesToolDefinition;
    private contextExtractor: ContextExtractor;

    constructor() {
        this.contextExtractor = new ContextExtractor();
    }

    /**
     * Get rich content preview for a note
     * This provides a better preview than the simple truncation in VectorSearchTool
     */
    private async getRichContentPreview(noteId: string, summarize: boolean): Promise<string> {
        try {
            const note = becca.getNote(noteId);
            if (!note) {
                return 'Note not found';
            }

            // Get the full content with proper formatting
            const formattedContent = await this.contextExtractor.getNoteContent(noteId);
            if (!formattedContent) {
                return 'No content available';
            }

            // If summarization is requested
            if (summarize) {
                // Try to get an LLM service for summarization
                try {
                    const llmService = await aiServiceManager.getService();
                    
                    const messages = [
                            {
                                role: "system" as const,
                                content: "Summarize the following note content concisely while preserving key information. Keep your summary to about 3-4 sentences."
                            },
                            {
                                role: "user" as const,
                                content: `Note title: ${note.title}\n\nContent:\n${formattedContent}`
                            }
                        ];

                        // Request summarization with safeguards to prevent recursion
                        const result = await llmService.generateChatCompletion(messages, {
                            temperature: 0.3,
                            maxTokens: 200,
                            // Type assertion to bypass type checking for special internal parameters
                            ...(({
                                bypassFormatter: true,
                                bypassContextProcessing: true
                            } as Record<string, boolean>))
                        });

                    if (result && result.text) {
                        return result.text;
                    }
                } catch (error) {
                    log.error(`Error summarizing content: ${error}`);
                    // Fall through to smart truncation if summarization fails
                }
            }

            try {
                // Fall back to smart truncation if summarization fails or isn't requested
                const previewLength = Math.min(formattedContent.length, 600);
                let preview = formattedContent.substring(0, previewLength);

                // Only add ellipsis if we've truncated the content
                if (previewLength < formattedContent.length) {
                    // Try to find a natural break point
                    const breakPoints = ['. ', '.\n', '\n\n', '\n', '. '];

                    for (const breakPoint of breakPoints) {
                        const lastBreak = preview.lastIndexOf(breakPoint);
                        if (lastBreak > previewLength * 0.6) { // At least 60% of the way through
                            preview = preview.substring(0, lastBreak + breakPoint.length);
                            break;
                        }
                    }

                    // Add ellipsis if truncated
                    preview += '...';
                }

                return preview;
            } catch (error) {
                log.error(`Error getting rich content preview: ${error}`);
                return 'Error retrieving content preview';
            }
        } catch (error) {
            log.error(`Error getting rich content preview: ${error}`);
            return 'Error retrieving content preview';
        }
    }

    /**
     * Extract keywords from a semantic query for alternative search suggestions
     */
    private extractKeywords(query: string): string {
        return query.split(' ')
            .filter(word => word.length > 3 && !['using', 'with', 'for', 'and', 'the', 'that', 'this'].includes(word.toLowerCase()))
            .slice(0, 3)
            .join(' ');
    }

    /**
     * Suggest broader search terms when specific searches fail
     */
    private suggestBroaderTerms(query: string): string {
        const broaderTermsMap: Record<string, string> = {
            'machine learning': 'AI technology',
            'productivity': 'work methods',
            'development': 'programming',
            'management': 'organization',
            'planning': 'strategy'
        };
        
        for (const [specific, broader] of Object.entries(broaderTermsMap)) {
            if (query.toLowerCase().includes(specific)) {
                return broader;
            }
        }
        
        // Default: take first significant word and make it broader
        const firstWord = query.split(' ').find(word => word.length > 3);
        return firstWord ? `${firstWord} concepts` : 'general topics';
    }

    /**
     * Execute the search notes tool
     */
    public async execute(args: {
        query: string,
        parentNoteId?: string,
        maxResults?: number,
        summarize?: boolean
    }): Promise<string | object> {
        try {
            const {
                query,
                parentNoteId,
                maxResults = 5,
                summarize = false
            } = args;

            log.info(`Executing search_notes tool - Query: "${query}", ParentNoteId: ${parentNoteId || 'not specified'}, MaxResults: ${maxResults}, Summarize: ${summarize}`);

            // Get the vector search tool from the AI service manager
            const vectorSearchTool = await getOrCreateVectorSearchTool();

            if (!vectorSearchTool) {
                return `Error: Vector search tool is not available. The system may still be initializing or there could be a configuration issue.`;
            }

            log.info(`Retrieved vector search tool from AI service manager`);

            // Check if searchNotes method exists
            if (!vectorSearchTool.searchNotes || typeof vectorSearchTool.searchNotes !== 'function') {
                log.error(`Vector search tool is missing searchNotes method`);
                return `Error: Vector search tool is improperly configured (missing searchNotes method).`;
            }

            // Execute the search
            log.info(`Performing semantic search for: "${query}"`);
            const searchStartTime = Date.now();
            const response = await vectorSearchTool.searchNotes(query, parentNoteId, maxResults);
            const results: Array<Record<string, unknown>> = response?.matches ?? [];
            const searchDuration = Date.now() - searchStartTime;

            log.info(`Search completed in ${searchDuration}ms, found ${results.length} matching notes`);

            if (results.length > 0) {
                // Log top results
                results.slice(0, 3).forEach((result: any, index: number) => {
                    log.info(`Result ${index + 1}: "${result.title}" (similarity: ${Math.round(result.similarity * 100)}%)`);
                });
            } else {
                log.info(`No matching notes found for query: "${query}"`);
            }

            // Get enhanced previews for each result
            const enhancedResults = await Promise.all(
                results.map(async (result: any) => {
                    const noteId = result.noteId;
                    const preview = await this.getRichContentPreview(noteId, summarize);

                    return {
                        noteId: noteId,
                        title: result?.title as string || '[Unknown title]',
                        preview: preview,
                        score: result?.score as number,
                        dateCreated: result?.dateCreated as string,
                        dateModified: result?.dateModified as string,
                        similarity: Math.round(result.similarity * 100) / 100,
                        parentId: result.parentId
                    };
                })
            );

            // Format the results with enhanced guidance
            if (results.length === 0) {
                return {
                    count: 0,
                    results: [],
                    query: query,
                    searchType: 'semantic',
                    message: 'No semantic matches found for your query.',
                    nextSteps: {
                        immediate: [
                            `Try keyword_search with specific terms: "${this.extractKeywords(query)}"`,
                            `Use attribute_search if looking for labeled/categorized notes`,
                            `Try broader search terms like "${this.suggestBroaderTerms(query)}"`
                        ],
                        tips: [
                            'Semantic search finds conceptual matches - try describing the topic differently',
                            'If you know specific words that appear in the notes, use keyword_search instead',
                            'Check if the content might be tagged with labels using attribute_search'
                        ]
                    }
                };
            } else {
                return {
                    count: enhancedResults.length,
                    results: enhancedResults,
                    query: query,
                    searchType: 'semantic',
                    message: 'Found semantic matches. Use noteId values with other tools.',
                    nextSteps: {
                        examine: `Use read_note with any noteId (e.g., "${enhancedResults[0].noteId}") to get full content`,
                        refine: parentNoteId ? 'Remove parentNoteId to search all notes' : `Add parentNoteId: "${enhancedResults[0].noteId}" to search within the first result's children`,
                        related: 'Search for related concepts or use different descriptive terms'
                    }
                };
            }
        } catch (error: unknown) {
            const errorMessage = error instanceof Error ? error.message : String(error);
            log.error(`Error executing search_notes tool: ${errorMessage}`);
            return `Error: ${errorMessage}`;
        }
    }
}