import type { Message, ChatCompletionOptions } from './ai_interface.js'; import aiServiceManager from './ai_service_manager.js'; import chatStorageService from './chat_storage_service.js'; export interface ChatSession { id: string; title: string; messages: Message[]; isStreaming?: boolean; options?: ChatCompletionOptions; } /** * Service for managing chat interactions and history */ export class ChatService { private activeSessions: Map = new Map(); private streamingCallbacks: Map void> = new Map(); /** * Create a new chat session */ async createSession(title?: string, initialMessages: Message[] = []): Promise { const chat = await chatStorageService.createChat(title || 'New Chat', initialMessages); const session: ChatSession = { id: chat.id, title: chat.title, messages: chat.messages, isStreaming: false }; this.activeSessions.set(chat.id, session); return session; } /** * Get an existing session or create a new one */ async getOrCreateSession(sessionId?: string): Promise { if (sessionId) { const existingSession = this.activeSessions.get(sessionId); if (existingSession) { return existingSession; } const chat = await chatStorageService.getChat(sessionId); if (chat) { const session: ChatSession = { id: chat.id, title: chat.title, messages: chat.messages, isStreaming: false }; this.activeSessions.set(chat.id, session); return session; } } return this.createSession(); } /** * Send a message in a chat session and get the AI response */ async sendMessage( sessionId: string, content: string, options?: ChatCompletionOptions, streamCallback?: (content: string, isDone: boolean) => void ): Promise { const session = await this.getOrCreateSession(sessionId); // Add user message const userMessage: Message = { role: 'user', content }; session.messages.push(userMessage); session.isStreaming = true; // Set up streaming if callback provided if (streamCallback) { this.streamingCallbacks.set(session.id, streamCallback); } try { // Immediately save the user message await chatStorageService.updateChat(session.id, session.messages); // Generate AI response const response = await aiServiceManager.generateChatCompletion( session.messages, options || session.options ); // Add assistant message const assistantMessage: Message = { role: 'assistant', content: response.text }; session.messages.push(assistantMessage); session.isStreaming = false; // Save the complete conversation await chatStorageService.updateChat(session.id, session.messages); // If first message, update the title based on content if (session.messages.length <= 2 && !session.title) { // Extract a title from the conversation const title = this.generateTitleFromMessages(session.messages); session.title = title; await chatStorageService.updateChat(session.id, session.messages, title); } // Notify streaming is complete if (streamCallback) { streamCallback(response.text, true); this.streamingCallbacks.delete(session.id); } return session; } catch (error: any) { session.isStreaming = false; console.error('Error in AI chat:', error); // Add error message so user knows something went wrong const errorMessage: Message = { role: 'assistant', content: `Error: Failed to generate response. ${error.message || 'Please check AI settings and try again.'}` }; session.messages.push(errorMessage); // Save the conversation with error await chatStorageService.updateChat(session.id, session.messages); // Notify streaming is complete with error if (streamCallback) { streamCallback(errorMessage.content, true); this.streamingCallbacks.delete(session.id); } return session; } } /** * Add context from the current note to the chat * * @param sessionId - The ID of the chat session * @param noteId - The ID of the note to add context from * @param useSmartContext - Whether to use smart context extraction (default: true) * @returns The updated chat session */ async addNoteContext(sessionId: string, noteId: string, useSmartContext = true): Promise { const session = await this.getOrCreateSession(sessionId); // Get the last user message to use as context for semantic search const lastUserMessage = [...session.messages].reverse() .find(msg => msg.role === 'user' && msg.content.length > 10)?.content || ''; let context; if (useSmartContext && lastUserMessage) { // Use smart context that considers the query for better relevance context = await contextExtractor.getSmartContext(noteId, lastUserMessage); } else { // Fall back to full context if smart context is disabled or no query available context = await contextExtractor.getFullContext(noteId); } const contextMessage: Message = { role: 'user', content: `Here is the content of the note I want to discuss:\n\n${context}\n\nPlease help me with this information.` }; session.messages.push(contextMessage); await chatStorageService.updateChat(session.id, session.messages); return session; } /** * Add semantically relevant context from a note based on a specific query * * @param sessionId - The ID of the chat session * @param noteId - The ID of the note to add context from * @param query - The specific query to find relevant information for * @returns The updated chat session */ async addSemanticNoteContext(sessionId: string, noteId: string, query: string): Promise { const session = await this.getOrCreateSession(sessionId); // Use semantic context that considers the query for better relevance const contextService = aiServiceManager.getContextService(); const context = await contextService.getSemanticContext(noteId, query); const contextMessage: Message = { role: 'user', content: `Here is the relevant information from my notes based on my query "${query}":\n\n${context}\n\nPlease help me understand this information in relation to my query.` }; session.messages.push(contextMessage); await chatStorageService.updateChat(session.id, session.messages); return session; } /** * Send a message with enhanced semantic note context */ async sendContextAwareMessage( sessionId: string, content: string, noteId: string, options?: ChatCompletionOptions, streamCallback?: (content: string, isDone: boolean) => void ): Promise { const session = await this.getOrCreateSession(sessionId); // Add user message const userMessage: Message = { role: 'user', content }; session.messages.push(userMessage); session.isStreaming = true; // Set up streaming if callback provided if (streamCallback) { this.streamingCallbacks.set(session.id, streamCallback); } try { // Immediately save the user message await chatStorageService.updateChat(session.id, session.messages); // Get the Trilium Context Service for enhanced context const contextService = aiServiceManager.getContextService(); // Get showThinking option if it exists const showThinking = options?.showThinking === true; // Get enhanced context for this note and query const enhancedContext = await contextService.getAgentToolsContext( noteId, content, showThinking ); // Prepend a system message with context const systemMessage: Message = { role: 'system', content: `You are an AI assistant helping with Trilium Notes. Use this context to answer the user's question:\n\n${enhancedContext}` }; // Create messages array with system message const messagesWithContext = [systemMessage, ...session.messages]; // Generate AI response const response = await aiServiceManager.generateChatCompletion( messagesWithContext, options ); // Add assistant message const assistantMessage: Message = { role: 'assistant', content: response.text }; session.messages.push(assistantMessage); session.isStreaming = false; // Save the complete conversation (without system message) await chatStorageService.updateChat(session.id, session.messages); // If first message, update the title if (session.messages.length <= 2 && (!session.title || session.title === 'New Chat')) { const title = this.generateTitleFromMessages(session.messages); session.title = title; await chatStorageService.updateChat(session.id, session.messages, title); } // Notify streaming is complete if (streamCallback) { streamCallback(response.text, true); this.streamingCallbacks.delete(session.id); } return session; } catch (error: any) { session.isStreaming = false; console.error('Error in context-aware chat:', error); // Add error message const errorMessage: Message = { role: 'assistant', content: `Error: Failed to generate response with note context. ${error.message || 'Please try again.'}` }; session.messages.push(errorMessage); // Save the conversation with error await chatStorageService.updateChat(session.id, session.messages); // Notify streaming is complete with error if (streamCallback) { streamCallback(errorMessage.content, true); this.streamingCallbacks.delete(session.id); } return session; } } /** * Get all user's chat sessions */ async getAllSessions(): Promise { const chats = await chatStorageService.getAllChats(); return chats.map(chat => ({ id: chat.id, title: chat.title, messages: chat.messages, isStreaming: this.activeSessions.get(chat.id)?.isStreaming || false })); } /** * Delete a chat session */ async deleteSession(sessionId: string): Promise { this.activeSessions.delete(sessionId); this.streamingCallbacks.delete(sessionId); return chatStorageService.deleteChat(sessionId); } /** * Generate a title from the first messages in a conversation */ private generateTitleFromMessages(messages: Message[]): string { if (messages.length < 2) { return 'New Chat'; } // Get the first user message const firstUserMessage = messages.find(m => m.role === 'user'); if (!firstUserMessage) { return 'New Chat'; } // Extract first line or first few words const firstLine = firstUserMessage.content.split('\n')[0].trim(); if (firstLine.length <= 30) { return firstLine; } // Take first 30 chars if too long return firstLine.substring(0, 27) + '...'; } } // Singleton instance const chatService = new ChatService(); export default chatService;