diff --git a/src/cache.js b/src/cache.js new file mode 100644 index 0000000..c8ad2e5 --- /dev/null +++ b/src/cache.js @@ -0,0 +1,2 @@ +"use strict"; +Object.defineProperty(exports, "__esModule", { value: true }); diff --git a/src/chat-view.js b/src/chat-view.js index 7d352a9..8ec069f 100644 --- a/src/chat-view.js +++ b/src/chat-view.js @@ -1,452 +1,457 @@ -'use strict'; -Object.defineProperty(exports, '__esModule', { value: true }); -exports.ChatView = void 0; -const obsidian_1 = require('obsidian'); -const MAX_STREAM_CHUNKS = 1000; -const ollama_client_1 = require('./ollama-client'); -const vault_indexer_1 = require('./vault-indexer'); -const tool_executor_1 = require('./tool-executor'); -const error_handler_1 = require('./error-handler'); +"use strict"; +Object.defineProperty(exports, "__esModule", { value: true }); +exports.ChatView = exports.VIEW_TYPE_OLLAMA_CHAT = void 0; +const obsidian_1 = require("obsidian"); +const ollama_client_1 = require("./ollama-client"); +const vault_indexer_1 = require("./vault-indexer"); +const tool_executor_1 = require("./tool-executor"); +const conversation_state_1 = require("./conversation-state"); +const error_handler_1 = require("./error-handler"); +exports.VIEW_TYPE_OLLAMA_CHAT = 'ollama-chat-view'; class ChatView extends obsidian_1.ItemView { - // Getters for testing - getSendButtonClickHandler() { - return this.sendButtonClickHandler; - } - getInputKeyDownHandler() { - return this.inputKeyDownHandler; - } - getNewChatButtonClickHandler() { - return this.newChatButtonClickHandler; - } - constructor(leaf, settings) { - super(leaf); - this.messages = []; - this.lastMessageEl = null; - this.newChatButton = null; - this.sendButton = null; - this.inputEl = null; - this.chatContainer = null; - this.sendButtonClickHandler = null; - this.inputKeyDownHandler = null; - this.newChatButtonClickHandler = null; - this.sendButtonClickWrapper = null; - this.inputKeyDownWrapper = null; - this.newChatButtonClickWrapper = null; - this.listenersAttached = false; - this.settings = settings; - this.ollamaClient = new ollama_client_1.OllamaClient( - settings.ollamaUrl, - settings.model, - undefined, - settings.cacheConfig - ); - this.vaultIndexer = new vault_indexer_1.VaultIndexer(this.app.vault); - this.toolExecutor = new tool_executor_1.ToolExecutor(this.app.vault, this.app); - } - updateSettings(newSettings) { - this.settings = newSettings; - this.ollamaClient = new ollama_client_1.OllamaClient( - newSettings.ollamaUrl, - newSettings.model, - undefined, - newSettings.cacheConfig - ); - void this.ollamaClient.initializeCache().catch(() => { - new obsidian_1.Notice( - 'Semantic cache unavailable: could not connect to ChromaDB. Check the ChromaDB URL in settings.' - ); - }); - } - async clearCache() { - await this.ollamaClient.clearCache(); - } - getViewType() { - return 'ollama-chat-view'; - } - getDisplayText() { - return 'Ollama Chat'; - } - async onOpen() { - try { - await this.ollamaClient.initializeCache(); - } catch { - new obsidian_1.Notice( - 'Semantic cache unavailable: could not connect to ChromaDB. Check the ChromaDB URL in settings.' - ); + // Getters for testing + getSendButtonClickHandler() { + return this.sendButtonClickHandler; } - this.render(); - this.removeEventListeners(); // Clean up any existing listeners before reattaching - this.setupEventListeners(); - } - onSettingsChange(newSettings) { - this.updateSettings(newSettings); - } - async onClose() { - this.ollamaClient.cancelStream(); - this.removeEventListeners(); - this.cleanupStreamingResources(); - this.lastMessageEl = null; - this.sendButton = null; - this.inputEl = null; - this.chatContainer = null; - return Promise.resolve(); - } - cleanupStreamingResources() { - // Only cleanup if there's still an active streaming message - const streamingMessage = this.messages.find((msg) => msg.isStreaming); - if (streamingMessage && this.lastMessageEl && this.lastMessageEl.parentElement) { - this.lastMessageEl.parentElement.removeChild(this.lastMessageEl); - this.lastMessageEl = null; + getInputKeyDownHandler() { + return this.inputKeyDownHandler; } - } - render() { - const container = - this.chatContainer || this.contentEl.createEl('div', { cls: 'ollama-chat-container' }); - this.chatContainer = container; - const inputContainer = - this.contentEl.querySelector('.ollama-input-container') || - this.contentEl.createEl('div', { cls: 'ollama-input-container' }); - const newChatContainer = - this.contentEl.querySelector('.ollama-new-chat-container') || - this.contentEl.createEl('div', { cls: 'ollama-new-chat-container' }); - const messagesSnapshot = [...this.messages]; - const nonStreamingMessages = messagesSnapshot.filter((msg) => !msg.isStreaming); - const existingMessages = container.querySelectorAll('.ollama-message'); - // Remove messages that are no longer in the array - for (const el of Array.from(existingMessages)) { - const id = el.getAttribute('data-msg-id'); - if (!id || !nonStreamingMessages.some((m) => m.id === id)) { - el.remove(); - } + getNewChatButtonClickHandler() { + return this.newChatButtonClickHandler; } - // Render non-streaming messages - for (const msg of nonStreamingMessages) { - const existingEl = container.querySelector(`.ollama-message[data-msg-id="${msg.id}"]`); - if (existingEl) { - existingEl.querySelector('.ollama-message-content').textContent = msg.content; - } else { - const messageEl = container.createEl('div', { cls: 'ollama-message' }); - messageEl.setAttribute('data-msg-id', msg.id); - messageEl.createEl('div', { cls: 'ollama-message-role', text: msg.role }); - const contentEl = messageEl.createEl('div', { cls: 'ollama-message-content' }); - contentEl.textContent = msg.content; - } + constructor(leaf, settings) { + super(leaf); + // State + this.messages = []; + this.lastMessageEl = null; + this.newChatButton = null; + this.sendButton = null; + this.inputEl = null; + this.chatContainer = null; + this.sendButtonClickHandler = null; + this.inputKeyDownHandler = null; + this.newChatButtonClickHandler = null; + this.sendButtonClickWrapper = null; + this.inputKeyDownWrapper = null; + this.newChatButtonClickWrapper = null; + this.listenersAttached = false; + this.messages = []; + this.lastMessageEl = null; + this.newChatButton = null; + this.sendButton = null; + this.inputEl = null; + this.chatContainer = null; + this.sendButtonClickHandler = null; + this.inputKeyDownHandler = null; + this.newChatButtonClickHandler = null; + this.sendButtonClickWrapper = null; + this.inputKeyDownWrapper = null; + this.newChatButtonClickWrapper = null; + this.listenersAttached = false; + this.settings = settings; + this.ollamaClient = new ollama_client_1.OllamaClient(settings.ollamaUrl, settings.model, undefined, settings.cacheConfig); + this.vaultIndexer = new vault_indexer_1.VaultIndexer(this.app.vault); + this.toolExecutor = new tool_executor_1.ToolExecutor(this.app.vault, this.app); + this.conversationStateManager = new conversation_state_1.ConversationStateManager(); } - // Re-attach streaming message if it exists - const streamingMessage = messagesSnapshot.find((msg) => msg.isStreaming); - if (streamingMessage && this.lastMessageEl) { - const existingStreamingEl = container.querySelector( - `.ollama-message[data-msg-id="${streamingMessage.id}"]` - ); - if (!existingStreamingEl) { - container.appendChild(this.lastMessageEl); - } + updateSettings(newSettings) { + this.settings = newSettings; + this.ollamaClient = new ollama_client_1.OllamaClient(newSettings.ollamaUrl, newSettings.model, undefined, newSettings.cacheConfig); + void this.ollamaClient.initializeCache().catch(() => { + new obsidian_1.Notice('Semantic cache unavailable: could not connect to ChromaDB. Check the ChromaDB URL in settings.'); + }); } - // Setup new chat button - if (!this.newChatButton) { - this.newChatButton = newChatContainer.createEl('button', { - cls: 'ollama-new-chat-button', - text: 'New Chat', - }); - this.newChatButtonClickHandler = this.getNewChatButtonClickHandler(); - this.newChatButton.addEventListener('click', this.newChatButtonClickHandler); - } else { - newChatContainer.appendChild(this.newChatButton); + async clearCache() { + await this.ollamaClient.clearCache(); } - // Setup input area - if (!this.inputEl) { - this.inputEl = inputContainer.createEl('textarea', { - cls: 'ollama-input', - attr: { placeholder: 'Type your message...' }, - }); - this.inputKeyDownHandler = this.getInputKeyDownHandler(); - this.inputEl.addEventListener('keydown', this.inputKeyDownHandler); - } else { - inputContainer.appendChild(this.inputEl); + getViewType() { + return 'ollama-chat-view'; } - // Setup send button - if (!this.sendButton) { - this.sendButton = inputContainer.createEl('button', { - cls: 'ollama-send-button', - text: 'Send', - }); - this.sendButtonClickHandler = this.getSendButtonClickHandler(); - this.sendButton.addEventListener('click', this.sendButtonClickHandler); - } else { - inputContainer.appendChild(this.sendButton); + getDisplayText() { + return 'Ollama Chat'; } - // Append containers to contentEl - this.contentEl.appendChild(newChatContainer); - this.contentEl.appendChild(inputContainer); - this.contentEl.appendChild(container); - // Focus input on open - this.inputEl.focus(); - } - setupEventListeners() { - if (this.listenersAttached) { - return; - } - this.sendButtonClickHandler = () => { - void this.handleUserInput(); - }; - this.inputKeyDownHandler = (event) => { - if (event.key === 'Enter' && !event.shiftKey) { - event.preventDefault(); - void this.handleUserInput(); - } - }; - this.newChatButtonClickHandler = () => { - this.clearConversation(); - }; - if (this.sendButton) { - this.sendButton.addEventListener('click', this.sendButtonClickHandler); - } - if (this.inputEl) { - this.inputEl.addEventListener('keydown', this.inputKeyDownHandler); - } - if (this.newChatButton) { - this.newChatButton.addEventListener('click', this.newChatButtonClickHandler); - } - this.listenersAttached = true; - } - removeEventListeners() { - if (!this.listenersAttached) { - return; - } - if (this.sendButton) { - this.sendButton.removeEventListener('click', this.sendButtonClickHandler); - } - if (this.inputEl) { - this.inputEl.removeEventListener('keydown', this.inputKeyDownHandler); - } - if (this.newChatButton) { - this.newChatButton.removeEventListener('click', this.newChatButtonClickHandler); - } - this.listenersAttached = false; - } - clearConversation() { - this.messages = []; - this.render(); - } - updateMessageById(id, updates) { - const index = this.messages.findIndex((m) => m.id === id); - if (index !== -1) { - this.messages[index] = { ...this.messages[index], ...updates }; - this.render(); - } - } - updateLastMessage(updates) { - const streamingMessage = this.messages.find((msg) => msg.isStreaming); - if (streamingMessage) { - const index = this.messages.findIndex((msg) => msg.id === streamingMessage.id); - if (index !== -1) { - this.messages[index] = { ...this.messages[index], ...updates }; + async onOpen() { + try { + await this.ollamaClient.initializeCache(); + } + catch { + new obsidian_1.Notice('Semantic cache unavailable: could not connect to ChromaDB. Check the ChromaDB URL in settings.'); + } this.render(); - } + this.removeEventListeners(); // Clean up any existing listeners before reattaching + this.setupEventListeners(); } - } - getTools() { - return [ - { - type: 'function', - function: { - name: 'read_vault_file', - description: 'Reads the content of a file from the vault', - parameters: { - type: 'object', - properties: { - path: { - type: 'string', - description: 'The path to the file to read', - }, - content: { - type: 'string', - description: 'The content of the file to read', - }, + onSettingsChange(newSettings) { + this.updateSettings(newSettings); + } + async onClose() { + this.ollamaClient.cancelStream(); + this.removeEventListeners(); + this.cleanupStreamingResources(); + this.lastMessageEl = null; + this.sendButton = null; + this.inputEl = null; + this.chatContainer = null; + return Promise.resolve(); + } + cleanupStreamingResources() { + // Only cleanup if there's still an active streaming message + const streamingMessage = this.messages.find((msg) => msg.isStreaming); + if (streamingMessage && this.lastMessageEl && this.lastMessageEl.parentElement) { + this.lastMessageEl.parentElement.removeChild(this.lastMessageEl); + this.lastMessageEl = null; + } + } + render() { + const container = this.chatContainer || this.contentEl.createEl('div', { cls: 'ollama-chat-container' }); + this.chatContainer = container; + const inputContainer = this.contentEl.querySelector('.ollama-input-container') || + this.contentEl.createEl('div', { cls: 'ollama-input-container' }); + const newChatContainer = this.contentEl.querySelector('.ollama-new-chat-container') || + this.contentEl.createEl('div', { cls: 'ollama-new-chat-container' }); + const messagesSnapshot = [...this.messages]; + const nonStreamingMessages = messagesSnapshot.filter((msg) => !msg.isStreaming); + const existingMessages = container.querySelectorAll('.ollama-message'); + // Remove messages that are no longer in the array + for (const el of Array.from(existingMessages)) { + const id = el.getAttribute('data-msg-id'); + if (!id || !nonStreamingMessages.some((m) => m.id === id)) { + el.remove(); + } + } + // Render non-streaming messages + for (const msg of nonStreamingMessages) { + const existingEl = container.querySelector(`.ollama-message[data-msg-id="${msg.id}"]`); + if (existingEl) { + const contentEl = existingEl.querySelector('.ollama-message-content'); + if (contentEl) { + contentEl.textContent = msg.content; + } + } + else { + const messageEl = container.createEl('div', { cls: 'ollama-message' }); + messageEl.setAttribute('data-msg-id', msg.id); + messageEl.createEl('div', { cls: 'ollama-message-role', text: msg.role }); + const contentEl = messageEl.createEl('div', { cls: 'ollama-message-content' }); + contentEl.textContent = msg.content; + } + } + // Re-attach streaming message if it exists + const streamingMessage = messagesSnapshot.find((msg) => msg.isStreaming); + if (streamingMessage && this.lastMessageEl) { + const existingStreamingEl = container.querySelector(`.ollama-message[data-msg-id="${streamingMessage.id}"]`); + if (!existingStreamingEl) { + container.appendChild(this.lastMessageEl); + } + } + // Setup new chat button + if (!this.newChatButton) { + this.newChatButton = newChatContainer.createEl('button', { + cls: 'ollama-new-chat-button', + text: 'New Chat', + }); + } + else { + newChatContainer.appendChild(this.newChatButton); + } + // Setup input area + if (!this.inputEl) { + this.inputEl = inputContainer.createEl('textarea', { + cls: 'ollama-input', + attr: { placeholder: 'Type your message...' }, + }); + } + else { + inputContainer.appendChild(this.inputEl); + } + // Setup send button + if (!this.sendButton) { + this.sendButton = inputContainer.createEl('button', { + cls: 'ollama-send-button', + text: 'Send', + }); + } + else { + inputContainer.appendChild(this.sendButton); + } + // Append containers to contentEl + this.contentEl.appendChild(newChatContainer); + this.contentEl.appendChild(inputContainer); + this.contentEl.appendChild(container); + // Focus input on open + this.inputEl.focus(); + } + setupEventListeners() { + if (this.listenersAttached) { + return; + } + this.sendButtonClickHandler = () => { + void this.handleUserInput(this.inputEl?.value); + }; + this.inputKeyDownHandler = (event) => { + if (event.key === 'Enter' && !event.shiftKey) { + event.preventDefault(); + void this.handleUserInput(this.inputEl?.value); + } + }; + this.newChatButtonClickHandler = () => { + this.clearConversation(); + }; + if (this.sendButton && this.sendButtonClickHandler) { + this.sendButton.addEventListener('click', this.sendButtonClickHandler); + } + if (this.inputEl && this.inputKeyDownHandler) { + this.inputEl.addEventListener('keydown', this.inputKeyDownHandler); + } + if (this.newChatButton && this.newChatButtonClickHandler) { + this.newChatButton.addEventListener('click', this.newChatButtonClickHandler); + } + this.listenersAttached = true; + } + removeEventListeners() { + if (!this.listenersAttached) { + return; + } + if (this.sendButton && this.sendButtonClickHandler) { + this.sendButton.removeEventListener('click', this.sendButtonClickHandler); + } + if (this.inputEl && this.inputKeyDownHandler) { + this.inputEl.removeEventListener('keydown', this.inputKeyDownHandler); + } + if (this.newChatButton && this.newChatButtonClickHandler) { + this.newChatButton.removeEventListener('click', this.newChatButtonClickHandler); + } + this.listenersAttached = false; + } + clearConversation() { + this.messages = []; + this.conversationStateManager.clear(); + this.render(); + } + updateMessageById(id, updates) { + const index = this.messages.findIndex((m) => m.id === id); + if (index !== -1) { + this.messages[index] = { ...this.messages[index], ...updates }; + this.render(); + } + } + updateLastMessage(updates) { + const streamingMessage = this.messages.find((msg) => msg.isStreaming); + if (streamingMessage) { + const index = this.messages.findIndex((msg) => msg.id === streamingMessage.id); + if (index !== -1) { + this.messages[index] = { ...this.messages[index], ...updates }; + this.render(); + } + } + } + getTools() { + return [ + { + type: 'function', + function: { + name: 'read_vault_file', + description: 'Reads the content of a file from the vault', + parameters: { + type: 'object', + properties: { + path: { + type: 'string', + description: 'The path to the file to read', + }, + content: { + type: 'string', + description: 'The content of the file to read', + }, + }, + required: ['path'], + }, + }, }, - required: ['path'], - }, - }, - }, - { - type: 'function', - function: { - name: 'search_vault_files', - description: 'Searches for files in the vault that match a given query', - parameters: { - type: 'object', - properties: { - query: { - type: 'string', - description: 'The search query to use', - }, - limit: { - type: 'number', - description: 'The maximum number of results to return', - }, + { + type: 'function', + function: { + name: 'search_vault_files', + description: 'Searches for files in the vault that match a given query', + parameters: { + type: 'object', + properties: { + query: { + type: 'string', + description: 'The search query to use', + }, + limit: { + type: 'number', + description: 'The maximum number of results to return', + }, + }, + required: ['query'], + }, + }, }, - required: ['query'], - }, - }, - }, - ]; - } - buildMessages(userMessageWithContext, tools) { - const systemContent = `You are an assistant that can help answer questions using the contents of a vault. + ]; + } + buildMessages(userMessageContent, tools) { + const systemContent = `You are an assistant that can help answer questions using the contents of a vault. The user can ask questions about their vault contents, and you should provide helpful responses based on the files. When a user asks for information, try to find relevant files using the search_vault_files tool and read their contents with the read_vault_file tool. Only use the tools if you need to access vault content that is not already in the context.`; - const systemMessage = { - role: 'system', - content: systemContent, - }; - const userMessageWithContext = { - role: 'user', - content: userMessageWithContext, - }; - const messages = [systemMessage, userMessageWithContext]; - if (tools && tools.length > 0) { - messages.push({ - role: 'assistant', - content: 'I have access to the following tools to help answer your questions:', - tool_calls: tools, - }); - } - return messages; - } - async processToolCalls(toolCalls, messages, tools, fullResponse, assistantMessageId) { - const MAX_TOOL_CALLS = 5; - const settledResults = await Promise.allSettled( - toolCalls.map(async (toolCall) => { - try { - const toolResult = await this.toolExecutor.executeTool( - toolCall.function.name, - toolCall.function.arguments - ); - return { - status: 'fulfilled', - value: toolResult, - }; - } catch (error) { - return { - status: 'rejected', - reason: error, - }; + const systemMessage = { + role: 'system', + content: systemContent, + }; + const userMessage = { + role: 'user', + content: userMessageContent, + }; + const messages = [systemMessage, userMessage]; + if (tools && tools.length > 0) { + messages.push({ + role: 'assistant', + content: 'I have access to the following tools to help answer your questions:', + }); } - }) - ); - const toolResults = settledResults - .filter((result) => result.status === 'fulfilled') - .map((result) => result.value); - const followUpMessages = toolResults.map((result) => { - return { - role: 'tool', - content: JSON.stringify(result), - tool_call_id: result.id, - }; - }); - const followUp = { - role: 'assistant', - content: 'I have processed your request using the following tools. Here are the results:', - tool_calls: tools, - }; - if (followUpMessages.length > 0) { - const finalMessages = [...messages, followUp, ...followUpMessages]; - const stream = await this.ollamaClient.streamChat(finalMessages, { temperature: 0.5 }); - let fullResponse = ''; - for await (const chunk of stream) { - fullResponse += chunk.message.content; - this.updateLastMessage({ - content: fullResponse, - isStreaming: true, + return messages; + } + async processToolCalls(toolCalls, messages, tools, fullResponse, assistantMessageId) { + const toolResults = (await Promise.all(toolCalls.slice(0, MAX_TOOL_CALLS).map(async (toolCall) => { + try { + const toolResult = await this.toolExecutor.handleToolCall(toolCall); + return { ...toolResult, id: toolCall.id }; + } + catch (error) { + error_handler_1.ErrorHandler.handleError(error, 'ChatView.handleUserInput'); + return null; + } + }))).filter((result) => result !== null); + const followUpMessages = toolResults.map((result) => { + return { + role: 'tool', + content: JSON.stringify(result), + tool_call_id: result.id ?? '', + }; }); - } - this.updateMessageById(assistantMessageId, { - content: fullResponse, - isStreaming: false, - }); + const followUp = { + role: 'assistant', + content: 'I have processed your request using the following tools. Here are the results:', + tool_calls: toolCalls, + }; + if (followUpMessages.length > 0) { + const finalMessages = [...messages, followUp, ...followUpMessages]; + const response = await this.ollamaClient.chat(finalMessages, tools); + const finalResponse = response.content || fullResponse; + this.updateMessageById(assistantMessageId, { + content: finalResponse, + isStreaming: false, + }); + } } - } - async handleUserInput() { - const userMessage = this.inputEl.value.trim(); - if (!userMessage) { - return; - } - const entries = await this.vaultIndexer.getVaultEntries(); - let context = ''; - const MAX_CONTEXT_LENGTH = 2000; - const messages = this.messages; - const tools = this.getTools(); - const messageId = crypto.randomUUID(); - const userMessageId = `${messageId}-user`; - const assistantMessageId = `${messageId}-assistant`; - const userChatMessage = { - id: userMessageId, - role: 'user', - content: userMessage, - timestamp: Date.now(), - }; - const assistantMessage = { - id: assistantMessageId, - role: 'assistant', - content: '', - timestamp: Date.now(), - isStreaming: true, - }; - this.messages = [...this.messages, userChatMessage, assistantMessage]; - this.render(); - this.inputEl.value = ''; - // Add the assistant message to the DOM to enable streaming - this.lastMessageEl = this.chatContainer.querySelector( - `.ollama-message[data-msg-id="${assistantMessageId}"]` - ); - try { - const stream = await this.ollamaClient.streamChat(this.buildMessages(userMessage, tools), { - temperature: 0.5, - }); - let fullResponse = ''; - let toolCalls = []; - let chunkCount = 0; - for await (const chunk of stream) { - if (chunk.message.content) { - fullResponse += chunk.message.content; - this.updateLastMessage({ - content: fullResponse, + async handleUserInput(inputValue) { + const userMessage = (inputValue ?? this.inputEl?.value ?? '').trim(); + if (!userMessage) { + return; + } + const MAX_CONTEXT_LENGTH = 2000; + const tools = this.getTools(); + const messageId = crypto.randomUUID(); + const userMessageId = `${messageId}-user`; + const assistantMessageId = `${messageId}-assistant`; + const userChatMessage = { + id: userMessageId, + role: 'user', + content: userMessage, + timestamp: Date.now(), + }; + const assistantMessage = { + id: assistantMessageId, + role: 'assistant', + content: '', + timestamp: Date.now(), isStreaming: true, - }); + }; + const previousStreamingEl = this.lastMessageEl; + this.messages = [...this.messages, userChatMessage, assistantMessage]; + this.render(); + if (this.inputEl) { + this.inputEl.value = ''; } - if (chunk.message.tool_calls) { - toolCalls = [...toolCalls, ...chunk.message.tool_calls]; + // Add the assistant message to the DOM to enable streaming + this.lastMessageEl = + this.chatContainer?.querySelector(`.ollama-message[data-msg-id="${assistantMessageId}"]`) ?? + null; + if (!this.lastMessageEl && previousStreamingEl) { + previousStreamingEl.classList.add('ollama-message'); + previousStreamingEl.setAttribute('data-msg-id', assistantMessageId); + this.contentEl.appendChild(previousStreamingEl); + this.lastMessageEl = previousStreamingEl; } - chunkCount++; - if (chunkCount > MAX_STREAM_CHUNKS) { - break; + try { + const entries = await this.vaultIndexer.searchVault(userMessage, this.settings.vaultSearchLimit); + const context = entries + .map((entry) => `${entry.title}\n${entry.content}`) + .join('\n\n') + .slice(0, MAX_CONTEXT_LENGTH); + const userMessageWithContext = context + ? `Relevant vault context:\n${context}\n\nUser question:\n${userMessage}` + : userMessage; + // Get the complete messages array for the LLM with all context layers + const completeMessages = this.conversationStateManager.getCompleteMessages(userMessageWithContext); + const stream = this.ollamaClient.streamChat(completeMessages, tools); + let fullResponse = ''; + let toolCalls = []; + let chunkCount = 0; + for await (const chunk of stream) { + if (chunk.content) { + fullResponse += chunk.content; + this.updateLastMessage({ + content: fullResponse, + isStreaming: true, + }); + } + if (chunk.tool_calls) { + toolCalls = [...toolCalls, ...chunk.tool_calls]; + } + chunkCount++; + if (chunkCount > MAX_STREAM_CHUNKS) { + break; + } + } + // Process tool calls if any + if (toolCalls.length > 0) { + await this.processToolCalls(toolCalls, completeMessages, tools, fullResponse, assistantMessageId); + } + // Update assistant message immutably — only if no tool calls were processed + if (toolCalls.length === 0) { + this.updateMessageById(assistantMessageId, { + content: fullResponse, + isStreaming: false, + }); + } + // Update conversation state with the assistant's response + this.conversationStateManager.updateShortTermContext({ role: 'user', content: userMessage }); + this.conversationStateManager.updateShortTermContext({ + role: 'assistant', + content: fullResponse, + }); + // Limit conversation history to prevent memory issues + if (this.messages.length > this.settings.maxMessageHistory) { + this.messages = this.messages.slice(-this.settings.maxMessageHistory); + } + this.render(); + } + catch (error) { + error_handler_1.ErrorHandler.handleError(error, 'ChatView.handleUserInput'); + this.updateMessageById(assistantMessageId, { + content: 'An error occurred while processing your request.', + isStreaming: false, + }); + } + finally { + // Clean up streaming resources regardless of outcome + this.cleanupStreamingResources(); } - } - // Process tool calls if any - if (toolCalls.length > 0) { - await this.processToolCalls( - toolCalls, - this.buildMessages(userMessage, tools), - tools, - fullResponse, - assistantMessageId - ); - } - // Update assistant message immutably — only if no tool calls were processed - if (toolCalls.length === 0) { - this.updateMessageById(assistantMessageId, { - isStreaming: false, - }); - } - // Limit conversation history to prevent memory issues - if (this.messages.length > this.settings.maxMessageHistory) { - this.messages = this.messages.slice(-this.settings.maxMessageHistory); - } - this.render(); - } finally { - // Clean up streaming resources regardless of outcome - this.cleanupStreamingResources(); } - } } +exports.ChatView = ChatView; +const MAX_STREAM_CHUNKS = 1000; +const MAX_TOOL_CALLS = 5; diff --git a/src/chat-view.ts b/src/chat-view.ts index 07add57..ea6407e 100755 --- a/src/chat-view.ts +++ b/src/chat-view.ts @@ -2,7 +2,8 @@ import { ItemView, Notice, WorkspaceLeaf } from 'obsidian'; import { OllamaClient } from './ollama-client'; import { VaultIndexer } from './vault-indexer'; import { ToolExecutor } from './tool-executor'; -import { PluginSettings } from './types'; +import { PluginSettings, OllamaMessage, OllamaTool, OllamaToolCall, ChatMessage } from './types'; +import { ConversationStateManager } from './conversation-state'; import { ErrorHandler } from './error-handler'; export const VIEW_TYPE_OLLAMA_CHAT = 'ollama-chat-view'; @@ -45,6 +46,7 @@ export class ChatView extends ItemView { ); this.vaultIndexer = new VaultIndexer(this.app.vault); this.toolExecutor = new ToolExecutor(this.app.vault, this.app); + this.conversationStateManager = new ConversationStateManager(); } updateSettings(newSettings: PluginSettings) { @@ -138,7 +140,10 @@ export class ChatView extends ItemView { for (const msg of nonStreamingMessages) { const existingEl = container.querySelector(`.ollama-message[data-msg-id="${msg.id}"]`); if (existingEl) { - existingEl.querySelector('.ollama-message-content').textContent = msg.content; + const contentEl = existingEl.querySelector('.ollama-message-content'); + if (contentEl) { + contentEl.textContent = msg.content; + } } else { const messageEl = container.createEl('div', { cls: 'ollama-message' }); messageEl.setAttribute('data-msg-id', msg.id); @@ -165,8 +170,6 @@ export class ChatView extends ItemView { cls: 'ollama-new-chat-button', text: 'New Chat', }); - this.newChatButtonClickHandler = this.getNewChatButtonClickHandler(); - this.newChatButton.addEventListener('click', this.newChatButtonClickHandler); } else { newChatContainer.appendChild(this.newChatButton); } @@ -177,8 +180,6 @@ export class ChatView extends ItemView { cls: 'ollama-input', attr: { placeholder: 'Type your message...' }, }); - this.inputKeyDownHandler = this.getInputKeyDownHandler(); - this.inputEl.addEventListener('keydown', this.inputKeyDownHandler); } else { inputContainer.appendChild(this.inputEl); } @@ -189,8 +190,6 @@ export class ChatView extends ItemView { cls: 'ollama-send-button', text: 'Send', }); - this.sendButtonClickHandler = this.getSendButtonClickHandler(); - this.sendButton.addEventListener('click', this.sendButtonClickHandler); } else { inputContainer.appendChild(this.sendButton); } @@ -210,13 +209,13 @@ export class ChatView extends ItemView { } this.sendButtonClickHandler = () => { - void this.handleUserInput(); + void this.handleUserInput(this.inputEl?.value); }; this.inputKeyDownHandler = (event: KeyboardEvent) => { if (event.key === 'Enter' && !event.shiftKey) { event.preventDefault(); - void this.handleUserInput(); + void this.handleUserInput(this.inputEl?.value); } }; @@ -224,15 +223,15 @@ export class ChatView extends ItemView { this.clearConversation(); }; - if (this.sendButton) { + if (this.sendButton && this.sendButtonClickHandler) { this.sendButton.addEventListener('click', this.sendButtonClickHandler); } - if (this.inputEl) { + if (this.inputEl && this.inputKeyDownHandler) { this.inputEl.addEventListener('keydown', this.inputKeyDownHandler); } - if (this.newChatButton) { + if (this.newChatButton && this.newChatButtonClickHandler) { this.newChatButton.addEventListener('click', this.newChatButtonClickHandler); } @@ -244,15 +243,15 @@ export class ChatView extends ItemView { return; } - if (this.sendButton) { + if (this.sendButton && this.sendButtonClickHandler) { this.sendButton.removeEventListener('click', this.sendButtonClickHandler); } - if (this.inputEl) { + if (this.inputEl && this.inputKeyDownHandler) { this.inputEl.removeEventListener('keydown', this.inputKeyDownHandler); } - if (this.newChatButton) { + if (this.newChatButton && this.newChatButtonClickHandler) { this.newChatButton.removeEventListener('click', this.newChatButtonClickHandler); } @@ -261,6 +260,7 @@ export class ChatView extends ItemView { clearConversation(): void { this.messages = []; + this.conversationStateManager.clear(); this.render(); } @@ -330,28 +330,27 @@ export class ChatView extends ItemView { ]; } - buildMessages(userMessageWithContext: string, tools?: OllamaTool[]): OllamaMessage[] { + buildMessages(userMessageContent: string, tools?: OllamaTool[]): OllamaMessage[] { const systemContent = `You are an assistant that can help answer questions using the contents of a vault. The user can ask questions about their vault contents, and you should provide helpful responses based on the files. When a user asks for information, try to find relevant files using the search_vault_files tool and read their contents with the read_vault_file tool. Only use the tools if you need to access vault content that is not already in the context.`; - const systemMessage = { + const systemMessage: OllamaMessage = { role: 'system', content: systemContent, }; - const userMessageWithContext = { + const userMessage: OllamaMessage = { role: 'user', - content: userMessageWithContext, + content: userMessageContent, }; - const messages: OllamaMessage[] = [systemMessage, userMessageWithContext]; + const messages: OllamaMessage[] = [systemMessage, userMessage]; if (tools && tools.length > 0) { messages.push({ role: 'assistant', content: 'I have access to the following tools to help answer your questions:', - tool_calls: tools, }); } @@ -365,86 +364,65 @@ export class ChatView extends ItemView { fullResponse: string, assistantMessageId: string ): Promise { - const MAX_TOOL_CALLS = 5; - const settledResults = await Promise.allSettled( - toolCalls.map(async (toolCall) => { + const toolResults = ( + await Promise.all( + toolCalls.slice(0, MAX_TOOL_CALLS).map(async (toolCall) => { try { - const toolResult = await this.toolExecutor.executeTool( - toolCall.function.name, - toolCall.function.arguments - ); - return { - status: 'fulfilled', - value: toolResult, - }; + const toolResult = await this.toolExecutor.handleToolCall(toolCall); + return { ...toolResult, id: toolCall.id }; } catch (error) { - return { - status: 'rejected', - reason: error, - }; + ErrorHandler.handleError(error, 'ChatView.handleUserInput'); + return null; } }) - ); + ) + ).filter((result): result is NonNullable => result !== null); - const toolResults = settledResults - .filter((result) => result.status === 'fulfilled') - .map((result) => result.value); - - const followUpMessages = toolResults.map((result) => { + const followUpMessages: OllamaMessage[] = toolResults.map((result) => { return { role: 'tool', content: JSON.stringify(result), - tool_call_id: result.id, + tool_call_id: result.id ?? '', }; }); - const followUp = { + const followUp: OllamaMessage = { role: 'assistant', content: 'I have processed your request using the following tools. Here are the results:', - tool_calls: tools, + tool_calls: toolCalls, }; if (followUpMessages.length > 0) { const finalMessages = [...messages, followUp, ...followUpMessages]; - const stream = await this.ollamaClient.streamChat(finalMessages, { temperature: 0.5 }); - let fullResponse = ''; - for await (const chunk of stream) { - fullResponse += chunk.message.content; - this.updateLastMessage({ - content: fullResponse, - isStreaming: true, - }); - } + const response = await this.ollamaClient.chat(finalMessages, tools); + const finalResponse = response.content || fullResponse; this.updateMessageById(assistantMessageId, { - content: fullResponse, + content: finalResponse, isStreaming: false, }); } } - async handleUserInput(): Promise { - const userMessage = this.inputEl.value.trim(); + async handleUserInput(inputValue?: string): Promise { + const userMessage = (inputValue ?? this.inputEl?.value ?? '').trim(); if (!userMessage) { return; } - const entries = await this.vaultIndexer.getVaultEntries(); - let context = ''; const MAX_CONTEXT_LENGTH = 2000; - const messages = this.messages; const tools = this.getTools(); const messageId = crypto.randomUUID(); const userMessageId = `${messageId}-user`; const assistantMessageId = `${messageId}-assistant`; - const userChatMessage = { + const userChatMessage: ChatMessage = { id: userMessageId, role: 'user', content: userMessage, timestamp: Date.now(), }; - const assistantMessage = { + const assistantMessage: ChatMessage = { id: assistantMessageId, role: 'assistant', content: '', @@ -452,35 +430,54 @@ export class ChatView extends ItemView { isStreaming: true, }; + const previousStreamingEl = this.lastMessageEl; this.messages = [...this.messages, userChatMessage, assistantMessage]; this.render(); - this.inputEl.value = ''; + if (this.inputEl) { + this.inputEl.value = ''; + } // Add the assistant message to the DOM to enable streaming - this.lastMessageEl = this.chatContainer.querySelector( - `.ollama-message[data-msg-id="${assistantMessageId}"]` - ); + this.lastMessageEl = + this.chatContainer?.querySelector(`.ollama-message[data-msg-id="${assistantMessageId}"]`) ?? + null; + if (!this.lastMessageEl && previousStreamingEl) { + previousStreamingEl.classList.add('ollama-message'); + previousStreamingEl.setAttribute('data-msg-id', assistantMessageId); + this.contentEl.appendChild(previousStreamingEl); + this.lastMessageEl = previousStreamingEl; + } try { - const stream = await this.ollamaClient.streamChat(this.buildMessages(userMessage, tools), { - temperature: 0.5, - }); + const entries = await this.vaultIndexer.searchVault(userMessage, this.settings.vaultSearchLimit); + const context = entries + .map((entry) => `${entry.title}\n${entry.content}`) + .join('\n\n') + .slice(0, MAX_CONTEXT_LENGTH); + const userMessageWithContext = context + ? `Relevant vault context:\n${context}\n\nUser question:\n${userMessage}` + : userMessage; + + // Get the complete messages array for the LLM with all context layers + const completeMessages = this.conversationStateManager.getCompleteMessages(userMessageWithContext); + + const stream = this.ollamaClient.streamChat(completeMessages, tools); let fullResponse = ''; let toolCalls: OllamaToolCall[] = []; let chunkCount = 0; for await (const chunk of stream) { - if (chunk.message.content) { - fullResponse += chunk.message.content; + if (chunk.content) { + fullResponse += chunk.content; this.updateLastMessage({ content: fullResponse, isStreaming: true, }); } - if (chunk.message.tool_calls) { - toolCalls = [...toolCalls, ...chunk.message.tool_calls]; + if (chunk.tool_calls) { + toolCalls = [...toolCalls, ...chunk.tool_calls]; } chunkCount++; @@ -493,7 +490,7 @@ export class ChatView extends ItemView { if (toolCalls.length > 0) { await this.processToolCalls( toolCalls, - this.buildMessages(userMessage, tools), + completeMessages, tools, fullResponse, assistantMessageId @@ -503,16 +500,30 @@ export class ChatView extends ItemView { // Update assistant message immutably — only if no tool calls were processed if (toolCalls.length === 0) { this.updateMessageById(assistantMessageId, { + content: fullResponse, isStreaming: false, }); } + // Update conversation state with the assistant's response + this.conversationStateManager.updateShortTermContext({ role: 'user', content: userMessage }); + this.conversationStateManager.updateShortTermContext({ + role: 'assistant', + content: fullResponse, + }); + // Limit conversation history to prevent memory issues if (this.messages.length > this.settings.maxMessageHistory) { this.messages = this.messages.slice(-this.settings.maxMessageHistory); } this.render(); + } catch (error) { + ErrorHandler.handleError(error, 'ChatView.handleUserInput'); + this.updateMessageById(assistantMessageId, { + content: 'An error occurred while processing your request.', + isStreaming: false, + }); } finally { // Clean up streaming resources regardless of outcome this.cleanupStreamingResources(); @@ -537,16 +548,8 @@ export class ChatView extends ItemView { private ollamaClient: OllamaClient; private vaultIndexer: VaultIndexer; private toolExecutor: ToolExecutor; -} - -// Type definitions -interface ChatMessage { - id: string; - role: 'user' | 'assistant'; - content: string; - timestamp: number; - isStreaming?: boolean; - tool_calls?: OllamaToolCall[]; + private conversationStateManager: ConversationStateManager; } const MAX_STREAM_CHUNKS = 1000; +const MAX_TOOL_CALLS = 5; diff --git a/src/constants.js b/src/constants.js index 87fe449..baad4e1 100644 --- a/src/constants.js +++ b/src/constants.js @@ -1,18 +1,17 @@ -'use strict'; -// Default plugin settings -Object.defineProperty(exports, '__esModule', { value: true }); +"use strict"; +Object.defineProperty(exports, "__esModule", { value: true }); exports.DEFAULT_SETTINGS = void 0; exports.DEFAULT_SETTINGS = { - ollamaUrl: 'http://localhost:11434', - model: 'llama3', - vaultSearchLimit: 3, - maxMessageHistory: 50, - lastIndexTime: 0, - cacheConfig: { - enabled: false, - similarityThreshold: 0.85, - collectionName: 'ollama_semantic_cache', - embeddingModel: 'nomic-embed-text', - chromaURL: 'http://localhost:8000', - }, + ollamaUrl: 'http://localhost:11434', + model: 'llama3', + vaultSearchLimit: 3, + maxMessageHistory: 50, + lastIndexTime: 0, + cacheConfig: { + enabled: false, + similarityThreshold: 0.85, + collectionName: 'ollama_semantic_cache', + embeddingModel: 'nomic-embed-text', + chromaURL: 'http://localhost:8000', + }, }; diff --git a/src/conversation-state.js b/src/conversation-state.js new file mode 100644 index 0000000..bec90e2 --- /dev/null +++ b/src/conversation-state.js @@ -0,0 +1,145 @@ +"use strict"; +// src/conversation-state.ts +Object.defineProperty(exports, "__esModule", { value: true }); +exports.ConversationStateManager = void 0; +class ConversationStateManager { + constructor() { + this.shortTermContext = []; + this.mediumTermContext = []; + this.longTermContext = []; + this.maxShortTermTurns = 10; + this.maxMediumTermMessages = 20; + // Initialize with default system context + this.longTermContext = [ + { + role: 'system', + content: `You are an assistant that can help answer questions using the contents of a vault. + The user can ask questions about their vault contents, and you should provide helpful responses based on the files. + When a user asks for information, try to find relevant files using the search_vault_files tool and read their contents with the read_vault_file tool. + Only use the tools if you need to access vault content that is not already in the context.`, + }, + ]; + } + /** + * Updates the short-term context with a new message + * @param message The message to add to short-term context + */ + updateShortTermContext(message) { + // Add new message + this.shortTermContext.push(message); + // Limit to max turns + if (this.shortTermContext.length > this.maxShortTermTurns) { + this.shortTermContext = this.shortTermContext.slice(-this.maxShortTermTurns); + } + } + /** + * Updates the medium-term context with a new message + * @param message The message to add to medium-term context + */ + updateMediumTermContext(message) { + // Add new message + this.mediumTermContext.push(message); + // Limit to max messages + if (this.mediumTermContext.length > this.maxMediumTermMessages) { + this.mediumTermContext = this.mediumTermContext.slice(-this.maxMediumTermMessages); + } + } + /** + * Sets the user's persona or core knowledge as long-term context + * @param personaContent The persona or core knowledge content + */ + setPersona(personaContent) { + // Remove any existing persona messages + this.longTermContext = this.longTermContext.filter((msg) => msg.role !== 'system' || + !msg.content.includes('You are an assistant that can help answer questions using the contents of a vault')); + // Add the new persona + this.longTermContext.push({ + role: 'system', + content: personaContent, + }); + } + /** + * Gets the combined conversation context for the current turn + * @param userMessage The user's current message + * @returns Complete conversation context with all three layers + */ + getConversationContext(_userMessage) { + return { + shortTermContext: this.shortTermContext, + mediumTermContext: this.mediumTermContext, + longTermContext: this.longTermContext, + }; + } + /** + * Gets the complete messages array for sending to the LLM + * @param userMessage The user's current message + * @returns Complete message array for the LLM + */ + getCompleteMessages(userMessage) { + const userMessageWithContext = { + role: 'user', + content: userMessage, + }; + // Build messages in the proper order: + // 1. Long-term context (user persona, system instructions) + // 2. Medium-term context (session knowledge base query results) + // 3. Short-term context (last N turns) + // 4. Current user message + return [ + ...this.longTermContext, + ...this.mediumTermContext, + ...this.shortTermContext, + userMessageWithContext, + ]; + } + /** + * Clears all conversation context + */ + clear() { + this.shortTermContext = []; + this.mediumTermContext = []; + this.longTermContext = [ + { + role: 'system', + content: `You are an assistant that can help answer questions using the contents of a vault. + The user can ask questions about their vault contents, and you should provide helpful responses based on the files. + When a user asks for information, try to find relevant files using the search_vault_files tool and read their contents with the read_vault_file tool. + Only use the tools if you need to access vault content that is not already in the context.`, + }, + ]; + } + /** + * Sets the medium-term context from a knowledge base query result + * @param queryResult The result from a knowledge base query + */ + setMediumTermContextFromQuery(queryResult) { + // Clear previous medium-term context + this.mediumTermContext = []; + // Add the query result as context + if (queryResult.trim()) { + this.mediumTermContext.push({ + role: 'system', + content: `Knowledge base results for current query:\n${queryResult}`, + }); + } + } + /** + * Gets the current short-term context + */ + getShortTermContext() { + return [...this.shortTermContext]; + } + /** + * Gets the current medium-term context + */ + getMediumTermContext() { + return [...this.mediumTermContext]; + } + /** + * Gets the current long-term context + */ + getLongTermContext() { + return [...this.longTermContext]; + } +} +exports.ConversationStateManager = ConversationStateManager; diff --git a/src/conversation-state.ts b/src/conversation-state.ts new file mode 100644 index 0000000..ba8ef54 --- /dev/null +++ b/src/conversation-state.ts @@ -0,0 +1,171 @@ +// src/conversation-state.ts + +import type { OllamaMessage } from './types'; + +export interface ConversationState { + shortTermContext: OllamaMessage[]; + mediumTermContext: OllamaMessage[]; + longTermContext: OllamaMessage[]; +} + +export class ConversationStateManager { + private shortTermContext: OllamaMessage[] = []; + private mediumTermContext: OllamaMessage[] = []; + private longTermContext: OllamaMessage[] = []; + private maxShortTermTurns: number = 10; + private maxMediumTermMessages: number = 20; + + constructor() { + // Initialize with default system context + this.longTermContext = [ + { + role: 'system', + content: `You are an assistant that can help answer questions using the contents of a vault. + The user can ask questions about their vault contents, and you should provide helpful responses based on the files. + When a user asks for information, try to find relevant files using the search_vault_files tool and read their contents with the read_vault_file tool. + Only use the tools if you need to access vault content that is not already in the context.`, + }, + ]; + } + + /** + * Updates the short-term context with a new message + * @param message The message to add to short-term context + */ + updateShortTermContext(message: OllamaMessage): void { + // Add new message + this.shortTermContext.push(message); + + // Limit to max turns + if (this.shortTermContext.length > this.maxShortTermTurns) { + this.shortTermContext = this.shortTermContext.slice(-this.maxShortTermTurns); + } + } + + /** + * Updates the medium-term context with a new message + * @param message The message to add to medium-term context + */ + updateMediumTermContext(message: OllamaMessage): void { + // Add new message + this.mediumTermContext.push(message); + + // Limit to max messages + if (this.mediumTermContext.length > this.maxMediumTermMessages) { + this.mediumTermContext = this.mediumTermContext.slice(-this.maxMediumTermMessages); + } + } + + /** + * Sets the user's persona or core knowledge as long-term context + * @param personaContent The persona or core knowledge content + */ + setPersona(personaContent: string): void { + // Remove any existing persona messages + this.longTermContext = this.longTermContext.filter( + (msg) => + msg.role !== 'system' || + !msg.content.includes( + 'You are an assistant that can help answer questions using the contents of a vault' + ) + ); + + // Add the new persona + this.longTermContext.push({ + role: 'system', + content: personaContent, + }); + } + + /** + * Gets the combined conversation context for the current turn + * @param userMessage The user's current message + * @returns Complete conversation context with all three layers + */ + getConversationContext(_userMessage: string): ConversationState { + return { + shortTermContext: this.shortTermContext, + mediumTermContext: this.mediumTermContext, + longTermContext: this.longTermContext, + }; + } + + /** + * Gets the complete messages array for sending to the LLM + * @param userMessage The user's current message + * @returns Complete message array for the LLM + */ + getCompleteMessages(userMessage: string): OllamaMessage[] { + const userMessageWithContext: OllamaMessage = { + role: 'user', + content: userMessage, + }; + + // Build messages in the proper order: + // 1. Long-term context (user persona, system instructions) + // 2. Medium-term context (session knowledge base query results) + // 3. Short-term context (last N turns) + // 4. Current user message + return [ + ...this.longTermContext, + ...this.mediumTermContext, + ...this.shortTermContext, + userMessageWithContext, + ]; + } + + /** + * Clears all conversation context + */ + clear(): void { + this.shortTermContext = []; + this.mediumTermContext = []; + this.longTermContext = [ + { + role: 'system', + content: `You are an assistant that can help answer questions using the contents of a vault. + The user can ask questions about their vault contents, and you should provide helpful responses based on the files. + When a user asks for information, try to find relevant files using the search_vault_files tool and read their contents with the read_vault_file tool. + Only use the tools if you need to access vault content that is not already in the context.`, + }, + ]; + } + + /** + * Sets the medium-term context from a knowledge base query result + * @param queryResult The result from a knowledge base query + */ + setMediumTermContextFromQuery(queryResult: string): void { + // Clear previous medium-term context + this.mediumTermContext = []; + + // Add the query result as context + if (queryResult.trim()) { + this.mediumTermContext.push({ + role: 'system', + content: `Knowledge base results for current query:\n${queryResult}`, + }); + } + } + + /** + * Gets the current short-term context + */ + getShortTermContext(): OllamaMessage[] { + return [...this.shortTermContext]; + } + + /** + * Gets the current medium-term context + */ + getMediumTermContext(): OllamaMessage[] { + return [...this.mediumTermContext]; + } + + /** + * Gets the current long-term context + */ + getLongTermContext(): OllamaMessage[] { + return [...this.longTermContext]; + } +} diff --git a/src/error-handler.js b/src/error-handler.js index c281335..d179f08 100644 --- a/src/error-handler.js +++ b/src/error-handler.js @@ -85,7 +85,7 @@ class ErrorHandler { return new types_1.NetworkError(message, statusCode); } static createApiError(message, statusCode) { - return new types_1.ApiError(message, statusCode); + return new types_1.ApiError(message, statusCode ?? 500); } static createValidationError(message, field) { const details = field ? { field, message } : undefined; diff --git a/src/error-handler.ts b/src/error-handler.ts index c675fa9..804e95f 100644 --- a/src/error-handler.ts +++ b/src/error-handler.ts @@ -109,7 +109,7 @@ export class ErrorHandler { } static createApiError(message: string, statusCode?: number): ApiError { - return new ApiError(message, statusCode); + return new ApiError(message, statusCode ?? 500); } static createValidationError(message: string, field?: string): ValidationError { diff --git a/src/indexing-pipeline/extraction.js b/src/indexing-pipeline/extraction.js index df48969..4296abb 100644 --- a/src/indexing-pipeline/extraction.js +++ b/src/indexing-pipeline/extraction.js @@ -81,7 +81,9 @@ class ContentExtractor { .replace(/^#.*?$/gm, '') .replace(/```.*?```/gs, '') .replace(/`.*?`/g, '') - .replace(/\[.*?\]\(.*?\)/g, '') + .replace(/\[(.*?)\]\(.*?\)/g, '$1') + .replace(/\*\*(.*?)\*\*/g, '$1') + .replace(/\*(.*?)\*/g, '$1') .trim(); } } diff --git a/src/indexing-pipeline/extraction.ts b/src/indexing-pipeline/extraction.ts index a332e72..b49fc97 100644 --- a/src/indexing-pipeline/extraction.ts +++ b/src/indexing-pipeline/extraction.ts @@ -1,7 +1,6 @@ // src/indexing-pipeline/extraction.ts -// VaultFile interface is defined locally since it's not exported from types -interface VaultFile { +export interface VaultFile { basename: string; path: string; } @@ -103,7 +102,9 @@ export class ContentExtractor { .replace(/^#.*?$/gm, '') .replace(/```.*?```/gs, '') .replace(/`.*?`/g, '') - .replace(/\[.*?\]\(.*?\)/g, '') + .replace(/\[(.*?)\]\(.*?\)/g, '$1') + .replace(/\*\*(.*?)\*\*/g, '$1') + .replace(/\*(.*?)\*/g, '$1') .trim(); } } diff --git a/src/indexing-pipeline/normalization.ts b/src/indexing-pipeline/normalization.ts index baa6044..f570468 100644 --- a/src/indexing-pipeline/normalization.ts +++ b/src/indexing-pipeline/normalization.ts @@ -3,13 +3,6 @@ // Import ExtractedContent interface from extraction module import { ExtractedContent } from './extraction'; -interface TokenizedContent { - tokens: string[]; - headings: string[]; - frontmatter: Record; - firstParagraph?: string; -} - interface NormalizedContent { path: string; title: string; diff --git a/src/indexing-pipeline/pipeline.js b/src/indexing-pipeline/pipeline.js index 81243a0..394993e 100644 --- a/src/indexing-pipeline/pipeline.js +++ b/src/indexing-pipeline/pipeline.js @@ -17,8 +17,11 @@ class IndexingPipeline { /** * Processes a vault file through the entire pipeline */ - async processFile(file, content) { + processFile(file, content) { try { + if (!content.trim()) { + return null; + } // Extraction step const extracted = this.extractor.extractFromFile(file, content); // Normalization/Enrichment step @@ -31,30 +34,30 @@ class IndexingPipeline { score: 0, // Score will be calculated during search }; } - catch (error) { + catch { return null; } } /** * Processes multiple files in batches */ - async processFilesInBatches(files, fileContents, batchSize = 10) { + processFilesInBatches(files, fileContents, batchSize = 10) { const results = []; const seenPaths = new Set(); for (let i = 0; i < files.length; i += batchSize) { const batch = files.slice(i, i + batchSize); - const batchResults = await Promise.all(batch.map(async (file) => { + const batchResults = batch.map((file) => { const content = fileContents[file.path]; if (!content) { return null; } - const entry = await this.processFile(file, content); + const entry = this.processFile(file, content); if (entry && !seenPaths.has(entry.path)) { seenPaths.add(entry.path); return entry; } return null; - })); + }); const validResults = batchResults.filter((result) => result !== null); results.push(...validResults); } diff --git a/src/indexing-pipeline/pipeline.ts b/src/indexing-pipeline/pipeline.ts index bd1760f..2d60530 100644 --- a/src/indexing-pipeline/pipeline.ts +++ b/src/indexing-pipeline/pipeline.ts @@ -1,7 +1,7 @@ // src/indexing-pipeline/pipeline.ts import { VaultIndexEntry } from '../types'; -import { ContentExtractor } from './extraction'; +import { ContentExtractor, VaultFile } from './extraction'; import { ContentNormalizer } from './normalization'; import { ContentVectorizer } from './vectorization'; @@ -27,8 +27,12 @@ export class IndexingPipeline { /** * Processes a vault file through the entire pipeline */ - async processFile(file: any, content: string): Promise { + processFile(file: VaultFile, content: string): VaultIndexEntry | null { try { + if (!content.trim()) { + return null; + } + // Extraction step const extracted = this.extractor.extractFromFile(file, content); @@ -42,7 +46,7 @@ export class IndexingPipeline { content: this.extractor.extractRawText(content).substring(0, 500), score: 0, // Score will be calculated during search }; - } catch (error) { + } catch { return null; } } @@ -50,31 +54,29 @@ export class IndexingPipeline { /** * Processes multiple files in batches */ - async processFilesInBatches( - files: any[], + processFilesInBatches( + files: VaultFile[], fileContents: Record, batchSize: number = 10 - ): Promise { + ): VaultIndexEntry[] { const results: VaultIndexEntry[] = []; const seenPaths = new Set(); for (let i = 0; i < files.length; i += batchSize) { const batch = files.slice(i, i + batchSize); - const batchResults = await Promise.all( - batch.map(async (file) => { + const batchResults = batch.map((file) => { const content = fileContents[file.path]; if (!content) { return null; } - const entry = await this.processFile(file, content); + const entry = this.processFile(file, content); if (entry && !seenPaths.has(entry.path)) { seenPaths.add(entry.path); return entry; } return null; - }) - ); + }); const validResults = batchResults.filter( (result): result is NonNullable => result !== null diff --git a/src/indexing-pipeline/vectorization.js b/src/indexing-pipeline/vectorization.js index 0e054eb..a2cc5f3 100644 --- a/src/indexing-pipeline/vectorization.js +++ b/src/indexing-pipeline/vectorization.js @@ -2,6 +2,7 @@ // src/indexing-pipeline/vectorization.ts Object.defineProperty(exports, "__esModule", { value: true }); exports.ContentVectorizer = void 0; +const utils_1 = require("../utils"); /** * Vectorizes content chunks using Ollama embeddings */ @@ -29,14 +30,23 @@ class ContentVectorizer { throw new Error(`Embedding failed with status ${response.status}`); } const data = await response.json(); + if (!this.isEmbeddingResponse(data)) { + throw new Error('Invalid embedding response'); + } return data.embedding; } catch (error) { // Return empty array on failure to maintain compatibility - console.warn(`Failed to generate embedding: ${String(error)}`); + utils_1.Logger.warn(`Failed to generate embedding: ${String(error)}`, 'indexing-pipeline'); return []; } } + isEmbeddingResponse(data) { + return (typeof data === 'object' && + data !== null && + Array.isArray(data.embedding) && + data.embedding.every((value) => typeof value === 'number')); + } /** * Creates a prompt from content chunk for embedding */ diff --git a/src/indexing-pipeline/vectorization.ts b/src/indexing-pipeline/vectorization.ts index 40fecc4..3c7fbd4 100644 --- a/src/indexing-pipeline/vectorization.ts +++ b/src/indexing-pipeline/vectorization.ts @@ -1,6 +1,7 @@ // src/indexing-pipeline/vectorization.ts import { ContentChunk } from './normalization'; +import { Logger } from '../utils'; interface VectorizationConfig { model: string; @@ -41,15 +42,28 @@ export class ContentVectorizer { throw new Error(`Embedding failed with status ${response.status}`); } - const data = await response.json(); + const data: unknown = await response.json(); + if (!this.isEmbeddingResponse(data)) { + throw new Error('Invalid embedding response'); + } + return data.embedding; } catch (error) { // Return empty array on failure to maintain compatibility - console.warn(`Failed to generate embedding: ${String(error)}`); + Logger.warn(`Failed to generate embedding: ${String(error)}`, 'indexing-pipeline'); return []; } } + private isEmbeddingResponse(data: unknown): data is { embedding: number[] } { + return ( + typeof data === 'object' && + data !== null && + Array.isArray((data as { embedding?: unknown }).embedding) && + (data as { embedding: unknown[] }).embedding.every((value) => typeof value === 'number') + ); + } + /** * Creates a prompt from content chunk for embedding */ diff --git a/src/main.js b/src/main.js index 51342b6..9c00ce7 100644 --- a/src/main.js +++ b/src/main.js @@ -1,204 +1,198 @@ -'use strict'; -Object.defineProperty(exports, '__esModule', { value: true }); -exports.OllamaPlugin = exports.SETTINGS_TAB_ID = void 0; -const obsidian_1 = require('obsidian'); -const chat_view_1 = require('./chat-view'); -const constants_1 = require('./constants'); -const ollama_client_1 = require('./ollama-client'); -const semantic_cache_1 = require('./semantic-cache'); -const vault_indexer_1 = require('./vault-indexer'); -exports.SETTINGS_TAB_ID = 'ollama-settings'; +"use strict"; +Object.defineProperty(exports, "__esModule", { value: true }); +const obsidian_1 = require("obsidian"); +const chat_view_1 = require("./chat-view"); +const constants_1 = require("./constants"); +const semantic_cache_1 = require("./semantic-cache"); class OllamaPlugin extends obsidian_1.Plugin { - async onload() { - await this.loadSettings(); - // Register the chat view - this.registerView( - chat_view_1.VIEW_TYPE_OLLAMA_CHAT, - (leaf) => new chat_view_1.ChatView(leaf, this.settings) - ); - // Add a command to open the chat view - this.addCommand({ - id: 'open-ollama-chat', - name: 'Open Ollama Chat', - callback: () => { - this.activateChatView(); - }, - }); - // Add a command to clear the semantic cache - this.addCommand({ - id: 'clear-semantic-cache', - name: 'Clear Semantic Cache', - callback: async () => { - await this.clearSemanticCache(); - new obsidian_1.Notice('Semantic cache cleared.'); - }, - }); - // Add a settings tab - this.addSettingTab(new OllamaSettingTab(this.app, this)); - // Initialize the semantic cache - this.semanticCache = new semantic_cache_1.SemanticCache(this.settings.cacheConfig); - try { - await this.semanticCache.initialize(); - } catch (error) { - console.error('Failed to initialize semantic cache:', error); - new obsidian_1.Notice('Semantic cache initialization failed. Check console for details.'); + constructor() { + super(...arguments); + this.settings = constants_1.DEFAULT_SETTINGS; } - } - async onunload() { - this.unregisterView(chat_view_1.VIEW_TYPE_OLLAMA_CHAT); - } - async loadSettings() { - this.settings = Object.assign({}, constants_1.DEFAULT_SETTINGS, await this.loadData()); - } - async saveSettings() { - await this.saveData(this.settings); - } - async activateChatView() { - const existing = this.app.workspace.getLeavesOfType(chat_view_1.VIEW_TYPE_OLLAMA_CHAT); - if (existing.length > 0) { - this.app.workspace.revealLeaf(existing[0]); - } else { - await this.app.workspace.getRightLeaf(false).setViewState({ - type: chat_view_1.VIEW_TYPE_OLLAMA_CHAT, - active: true, - }); + async onload() { + await this.loadSettings(); + // Register the chat view + this.registerView('ollama-chat-view', (leaf) => new chat_view_1.ChatView(leaf, this.settings)); + // Add a command to open the chat view + this.addCommand({ + id: 'open-ollama-chat', + name: 'Open Ollama Chat', + callback: async () => { + await this.activateChatView(); + }, + }); + // Add a command to clear the semantic cache + this.addCommand({ + id: 'clear-semantic-cache', + name: 'Clear Semantic Cache', + callback: async () => { + await this.clearSemanticCache(); + new obsidian_1.Notice('Semantic cache cleared.'); + }, + }); + // Add a settings tab + this.addSettingTab(new OllamaSettingTab(this.app, this)); + // Initialize the semantic cache + if (this.settings.cacheConfig) { + this.semanticCache = new semantic_cache_1.SemanticCacheService(this.settings.ollamaUrl, this.settings.cacheConfig); + try { + await this.semanticCache.initialize(); + } + catch { + new obsidian_1.Notice('Semantic cache initialization failed. Check console for details.'); + } + } } - } - async clearSemanticCache() { - if (this.semanticCache) { - await this.semanticCache.clear(); + // eslint-disable-next-line @typescript-eslint/no-misused-promises + onunload() { + // Clean up any active semantic cache resources on plugin unload + // Using fire-and-forget pattern since onunload cannot be async per Obsidian API + if (this.semanticCache) { + void this.semanticCache.clearCache(); + } + // No explicit unregisterView needed; relying on Obsidian lifecycle management. + } + async loadSettings() { + const loadedSettings = ((await this.loadData()) ?? {}); + this.settings = Object.assign({}, constants_1.DEFAULT_SETTINGS, loadedSettings); + } + async saveSettings() { + await this.saveData(this.settings); + } + async activateChatView() { + const existing = this.app.workspace.getLeavesOfType('ollama-chat-view'); + if (existing.length > 0) { + await this.app.workspace.revealLeaf(existing[0]); + } + else { + const leaf = this.app.workspace.getRightLeaf(false); + if (leaf) { + await leaf.setViewState({ + type: 'ollama-chat-view', + active: true, + }); + } + } + } + async clearSemanticCache() { + if (this.semanticCache) { + await this.semanticCache.clearCache(); + } + } + notifyChatViews() { + const leaves = this.app.workspace.getLeavesOfType('ollama-chat-view'); + leaves.forEach((leaf) => { + if (leaf.view instanceof chat_view_1.ChatView) { + leaf.view.updateSettings(this.settings); + } + }); } - } - notifyChatViews() { - const leaves = this.app.workspace.getLeavesOfType(chat_view_1.VIEW_TYPE_OLLAMA_CHAT); - leaves.forEach((leaf) => { - if (leaf.view instanceof chat_view_1.ChatView) { - leaf.view.updateSettings(this.settings); - } - }); - } } -exports.OllamaPlugin = OllamaPlugin; +exports.default = OllamaPlugin; class OllamaSettingTab extends obsidian_1.PluginSettingTab { - constructor(app, plugin) { - super(app, plugin); - this.plugin = plugin; - } - display() { - const { containerEl } = this; - containerEl.empty(); - containerEl.createEl('h2', { text: 'Ollama Settings' }); - new obsidian_1.Setting(containerEl) - .setName('Ollama URL') - .setDesc('URL for your Ollama instance (default: http://localhost:11434)') - .addText((text) => - text.setValue(this.plugin.settings.ollamaUrl).onChange(async (value) => { - this.plugin.settings.ollamaUrl = value; - await this.plugin.saveSettings(); - }) - ); - new obsidian_1.Setting(containerEl) - .setName('Model') - .setDesc('Ollama model to use (default: llama3)') - .addText((text) => - text.setValue(this.plugin.settings.model).onChange(async (value) => { - this.plugin.settings.model = value; - await this.plugin.saveSettings(); - }) - ); - new obsidian_1.Setting(containerEl) - .setName('Vault Search Limit') - .setDesc('Maximum number of vault entries to include in context (default: 3)') - .addText((text) => - text.setValue(String(this.plugin.settings.vaultSearchLimit)).onChange(async (value) => { - const parsed = parseInt(value); - if (!isNaN(parsed) && parsed > 0) { - this.plugin.settings.vaultSearchLimit = parsed; + constructor(app, plugin) { + super(app, plugin); + this.plugin = plugin; + } + display() { + const { containerEl } = this; + containerEl.empty(); + containerEl.createEl('h2', { text: 'Ollama Settings' }); + new obsidian_1.Setting(containerEl) + .setName('Ollama URL') + .setDesc('URL for your Ollama instance (default: http://localhost:11434)') + .addText((text) => text.setValue(this.plugin.settings.ollamaUrl).onChange(async (value) => { + this.plugin.settings.ollamaUrl = value; await this.plugin.saveSettings(); - } else { - new obsidian_1.Notice('Vault search limit must be a positive integer.'); - } - }) - ); - new obsidian_1.Setting(containerEl) - .setName('Max Message History') - .setDesc('Maximum number of messages to keep in conversation history (default: 50)') - .addText((text) => - text.setValue(String(this.plugin.settings.maxMessageHistory)).onChange(async (value) => { - const parsed = parseInt(value); - if (!isNaN(parsed) && parsed > 0) { - this.plugin.settings.maxMessageHistory = parsed; + })); + new obsidian_1.Setting(containerEl) + .setName('Model') + .setDesc('Ollama model to use (default: llama3)') + .addText((text) => text.setValue(this.plugin.settings.model).onChange(async (value) => { + this.plugin.settings.model = value; await this.plugin.saveSettings(); - } else { - new obsidian_1.Notice('Max message history must be a positive integer.'); - } - }) - ); - new obsidian_1.Setting(containerEl) - .setName('Enable Semantic Cache') - .setDesc('Use semantic cache to store and retrieve previous responses') - .addToggle((toggle) => - toggle.setValue(this.plugin.settings.cacheConfig.enabled).onChange(async (value) => { - this.plugin.settings.cacheConfig.enabled = value; - await this.plugin.saveSettings(); - this.plugin.notifyChatViews(); - }) - ); - new obsidian_1.Setting(containerEl) - .setName('ChromaDB URL') - .setDesc('URL for your ChromaDB instance (default: http://localhost:8000)') - .addText((text) => - text - .setValue(this.plugin.settings.cacheConfig.chromaURL || 'http://localhost:8000') - .onChange(async (value) => { + })); + new obsidian_1.Setting(containerEl) + .setName('Vault Search Limit') + .setDesc('Maximum number of vault entries to include in context (default: 3)') + .addText((text) => text.setValue(String(this.plugin.settings.vaultSearchLimit)).onChange(async (value) => { + const parsed = parseInt(value); + if (!isNaN(parsed) && parsed > 0) { + this.plugin.settings.vaultSearchLimit = parsed; + await this.plugin.saveSettings(); + } + else { + new obsidian_1.Notice('Vault search limit must be a positive integer.'); + } + })); + new obsidian_1.Setting(containerEl) + .setName('Max Message History') + .setDesc('Maximum number of messages to keep in conversation history (default: 50)') + .addText((text) => text.setValue(String(this.plugin.settings.maxMessageHistory)).onChange(async (value) => { + const parsed = parseInt(value); + if (!isNaN(parsed) && parsed > 0) { + this.plugin.settings.maxMessageHistory = parsed; + await this.plugin.saveSettings(); + } + else { + new obsidian_1.Notice('Max message history must be a positive integer.'); + } + })); + new obsidian_1.Setting(containerEl) + .setName('Enable Semantic Cache') + .setDesc('Use semantic cache to store and retrieve previous responses') + .addToggle((toggle) => toggle.setValue(this.plugin.settings.cacheConfig.enabled).onChange(async (value) => { + this.plugin.settings.cacheConfig.enabled = value; + await this.plugin.saveSettings(); + this.plugin.notifyChatViews(); + })); + new obsidian_1.Setting(containerEl) + .setName('ChromaDB URL') + .setDesc('URL for your ChromaDB instance (default: http://localhost:8000)') + .addText((text) => text + .setValue(this.plugin.settings.cacheConfig.chromaURL || 'http://localhost:8000') + .onChange(async (value) => { this.plugin.settings.cacheConfig.chromaURL = value; await this.plugin.saveSettings(); - }) - ); - new obsidian_1.Setting(containerEl) - .setName('Cache Embedding Model') - .setDesc('Ollama model used to generate embeddings for the semantic cache') - .addText((text) => - text.setValue(this.plugin.settings.cacheConfig.embeddingModel).onChange(async (value) => { - this.plugin.settings.cacheConfig.embeddingModel = value; - await this.plugin.saveSettings(); - this.plugin.notifyChatViews(); - }) - ); - new obsidian_1.Setting(containerEl) - .setName('Cache Similarity Threshold') - .setDesc( - 'Minimum cosine similarity (0–1) for a cache hit. Higher values require closer matches.' - ) - .addText((text) => - text - .setValue(String(this.plugin.settings.cacheConfig.similarityThreshold)) - .onChange(async (value) => { + })); + new obsidian_1.Setting(containerEl) + .setName('Cache Embedding Model') + .setDesc('Ollama model used to generate embeddings for the semantic cache') + .addText((text) => text.setValue(this.plugin.settings.cacheConfig.embeddingModel).onChange(async (value) => { + this.plugin.settings.cacheConfig.embeddingModel = value; + await this.plugin.saveSettings(); + this.plugin.notifyChatViews(); + })); + new obsidian_1.Setting(containerEl) + .setName('Cache Similarity Threshold') + .setDesc('Minimum cosine similarity (0–1) for a cache hit. Higher values require closer matches.') + .addText((text) => text + .setValue(String(this.plugin.settings.cacheConfig.similarityThreshold)) + .onChange(async (value) => { const parsed = parseFloat(value); if (!isNaN(parsed) && parsed >= 0 && parsed <= 1) { - this.plugin.settings.cacheConfig.similarityThreshold = parsed; - await this.plugin.saveSettings(); - } else { - new obsidian_1.Notice('Similarity threshold must be a number between 0 and 1.'); + this.plugin.settings.cacheConfig.similarityThreshold = parsed; + await this.plugin.saveSettings(); } - }) - ); - new obsidian_1.Setting(containerEl) - .setName('Clear Semantic Cache') - .setDesc('Delete all cached responses from ChromaDB') - .addButton((button) => - button.setButtonText('Clear Cache').onClick(async () => { - try { - await this.plugin.clearSemanticCache(); - new obsidian_1.Notice('Semantic cache cleared.'); - } catch { - new obsidian_1.Notice('Failed to clear semantic cache. Is ChromaDB running?'); - } - }) - ); - } - hide() { - // Clear the container to prevent duplicate elements - this.containerEl.empty(); - } + else { + new obsidian_1.Notice('Similarity threshold must be a number between 0 and 1.'); + } + })); + new obsidian_1.Setting(containerEl) + .setName('Clear Semantic Cache') + .setDesc('Delete all cached responses from ChromaDB') + .addButton((button) => button.setButtonText('Clear Cache').onClick(async () => { + try { + await this.plugin.clearSemanticCache(); + new obsidian_1.Notice('Semantic cache cleared.'); + } + catch { + new obsidian_1.Notice('Failed to clear semantic cache. Is ChromaDB running?'); + } + })); + } + hide() { + // Clear the container to prevent duplicate elements + this.containerEl.empty(); + } } diff --git a/src/main.ts b/src/main.ts index 85b0356..b03ad3e 100755 --- a/src/main.ts +++ b/src/main.ts @@ -64,10 +64,7 @@ export default class OllamaPlugin extends Plugin { } async loadSettings() { - // Obsidian's loadData() returns any, which is unavoidable in this API - // @ts-expect-error - Obsidian's loadData() returns any - const loadedSettings = await this.loadData(); - // @ts-expect-error - Merging with DEFAULT_SETTINGS + const loadedSettings = ((await this.loadData()) ?? {}) as Partial; this.settings = Object.assign({}, DEFAULT_SETTINGS, loadedSettings); } @@ -78,7 +75,7 @@ export default class OllamaPlugin extends Plugin { async activateChatView() { const existing = this.app.workspace.getLeavesOfType('ollama-chat-view'); if (existing.length > 0) { - this.app.workspace.revealLeaf(existing[0]); + await this.app.workspace.revealLeaf(existing[0]); } else { const leaf = this.app.workspace.getRightLeaf(false); if (leaf) { diff --git a/src/ollama-client.js b/src/ollama-client.js index af6bb27..46a6d0c 100644 --- a/src/ollama-client.js +++ b/src/ollama-client.js @@ -1,179 +1,294 @@ -'use strict'; +"use strict"; // src/ollama-client.ts -Object.defineProperty(exports, '__esModule', { value: true }); +Object.defineProperty(exports, "__esModule", { value: true }); exports.OllamaClient = void 0; -const types_1 = require('./types'); -const utils_1 = require('./utils'); -const semantic_cache_1 = require('./semantic-cache'); +const types_1 = require("./types"); +const utils_1 = require("./utils"); +const semantic_cache_1 = require("./semantic-cache"); class OllamaClient { - constructor(baseURL, model, fetchFn, cacheConfig) { - this.maxRetries = 3; - this.currentStreamController = null; - this.baseURL = baseURL; - this.model = model; - this.fetchFn = fetchFn ?? fetch; - if (cacheConfig?.enabled) { - this.cacheService = new semantic_cache_1.SemanticCacheService(baseURL, cacheConfig); - } - } - async initializeCache() { - if (this.cacheService) { - await this.cacheService.initialize(); - } - } - async clearCache() { - if (this.cacheService) { - await this.cacheService.clearCache(); - } - } - cancelStream() { - if (this.currentStreamController) { - this.currentStreamController.abort(); - this.currentStreamController = null; - } - } - async *streamChat(messages, tools = []) { - // Bypass cache if tools are involved to prevent state corruption - if (tools.length > 0) { - yield* this.streamChatWithRetry(messages, tools, 0); - return; - } - // Find the last user message - const lastUserMsg = [...messages].reverse().find((m) => m.role === 'user'); - if (lastUserMsg && this.cacheService) { - const cached = await this.cacheService.getCache(lastUserMsg.content); - if (cached) { - yield { role: 'assistant', content: cached, tool_calls: [] }; - return; - } - } - const chunks = []; - for await (const chunk of this.streamChatWithRetry(messages, tools, 0)) { - chunks.push(chunk); - yield chunk; - } - const fullContent = chunks.map((c) => c.content).join(''); - if (this.cacheService && lastUserMsg) { - void this.cacheService.setCache(lastUserMsg.content, fullContent); - } - } - async chat(messages, tools = []) { - // Bypass cache if tools are involved to prevent state corruption - if (tools.length > 0) { - return this.chatWithRetry(messages, tools, 0); - } - // Find the last user message - const lastUserMsg = [...messages].reverse().find((m) => m.role === 'user'); - if (lastUserMsg && this.cacheService) { - const cached = await this.cacheService.getCache(lastUserMsg.content); - if (cached) { - return { content: cached }; - } - } - const response = await this.chatWithRetry(messages, tools, 0); - if (this.cacheService && lastUserMsg) { - void this.cacheService.setCache(lastUserMsg.content, response.content); - } - return response; - } - async *streamChatWithRetry(messages, tools, retryCount) { - const controller = new AbortController(); - this.currentStreamController = controller; - try { - const response = await this.fetchFn(`${this.baseURL}/api/chat`, { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - }, - body: JSON.stringify({ - model: this.model, - messages: messages, - stream: true, - tools: tools, - }), - signal: controller.signal, - }); - if (!response.ok) { - throw new Error(`HTTP error! status: ${response.status}`); - } - if (!response.body) { - throw new Error('Response body is null'); - } - const reader = response.body.getReader(); - const decoder = new TextDecoder(); - let buffer = ''; - while (true) { - const { done, value } = await reader.read(); - if (done) { - break; + constructor(baseURL, model, fetchFn, cacheConfig) { + this.maxRetries = 3; + this.maxMalformedChunks = 50; + this.currentStreamController = null; + this.baseURL = baseURL; + this.model = model; + this.fetchFn = fetchFn ?? fetch; + if (cacheConfig?.enabled) { + this.cacheService = new semantic_cache_1.SemanticCacheService(baseURL, cacheConfig); + void this.cacheService.initialize(); } - buffer += decoder.decode(value); - const lines = buffer.split('\n'); - buffer = lines.pop() || ''; - for (const line of lines) { - if (line.trim() === '') { - continue; - } - try { - const parsed = JSON.parse(line); - yield parsed.message; - } catch (error) { - // Log but don't throw - malformed JSON is not critical - console.error('Failed to parse chunk:', line, error); - } + } + async initializeCache() { + if (this.cacheService) { + await this.cacheService.initialize(); } - } - if (buffer.trim() !== '') { + } + async clearCache() { + if (this.cacheService) { + await this.cacheService.clearCache(); + } + } + cancelStream() { + if (this.currentStreamController) { + this.currentStreamController.abort(); + this.currentStreamController = null; + } + } + async *streamChat(messages, tools = []) { + // Bypass cache if tools are involved to prevent state corruption + if (tools.length > 0) { + yield* this.streamChatWithRetry(messages, tools, 0); + return; + } + // Find the last user message + const lastUserMsg = [...messages].reverse().find((m) => m.role === 'user'); + if (lastUserMsg && this.cacheService) { + const cached = await this.cacheService.getCache(lastUserMsg.content); + if (cached) { + yield { role: 'assistant', content: cached, tool_calls: [] }; + return; + } + } + const chunks = []; + for await (const chunk of this.streamChatWithRetry(messages, tools, 0)) { + chunks.push(chunk); + yield chunk; + } + const fullContent = chunks.map((c) => c.content).join(''); + if (this.cacheService && lastUserMsg) { + void this.cacheService.setCache(lastUserMsg.content, fullContent); + } + } + async chat(messages, tools = []) { + // Bypass cache if tools are involved to prevent state corruption + if (tools.length > 0) { + return this.chatWithRetry(messages, tools, 0); + } + // Find the last user message + const lastUserMsg = [...messages].reverse().find((m) => m.role === 'user'); + if (lastUserMsg && this.cacheService) { + const cached = await this.cacheService.getCache(lastUserMsg.content); + if (cached) { + return { role: 'assistant', content: cached }; + } + } + const response = await this.chatWithRetry(messages, tools, 0); + if (this.cacheService && lastUserMsg) { + void this.cacheService.setCache(lastUserMsg.content, response.content); + } + return response; + } + async streamChatAsPromise(messages, tools = []) { + let content = ''; + let role = 'assistant'; + let toolCalls; + for await (const chunk of this.streamChat(messages, tools)) { + role = chunk.role ?? role; + content += chunk.content ?? ''; + if (chunk.tool_calls) { + toolCalls = [...(toolCalls ?? []), ...chunk.tool_calls]; + } + } + return { role, content, tool_calls: toolCalls }; + } + async *streamChatWithRetry(messages, tools = [], retryCount) { + const controller = new AbortController(); + this.currentStreamController = controller; + let reader = null; try { - const parsed = JSON.parse(buffer); - yield parsed.message; - } catch (error) { - console.error('Failed to parse final chunk:', buffer, error); + const response = await this.fetchFn(`${this.baseURL}/api/chat`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + model: this.model, + messages: messages, + tools: tools, + stream: true, + }), + signal: controller.signal, + }); + if (!response.ok) { + throw new types_1.ApiError(`Ollama API error: ${response.status}`, response.status); + } + if (!response.body) { + throw new Error('No response body'); + } + const contentType = response.headers?.get?.('content-type'); + if (contentType && !contentType.includes('application/x-ndjson')) { + throw new Error('Invalid response format'); + } + reader = response.body.getReader(); + const decoder = new TextDecoder(); + let buffer = ''; + let malformedChunks = 0; + while (true) { + const { done, value } = await reader.read(); + if (done) { + break; + } + buffer += decoder.decode(value); + const lines = buffer.split('\n'); + buffer = lines.pop() || ''; + for (const line of lines) { + if (line.trim() === '') { + continue; + } + let parsed; + try { + parsed = this.parseChatResponse(line); + } + catch (error) { + malformedChunks++; + const errorMessage = error instanceof Error ? error.message : String(error); + utils_1.Logger.warn(`Skipped malformed chunk: ${line.slice(0, 50)}... - ${errorMessage}`, 'ollama-client'); + if (malformedChunks > this.maxMalformedChunks) { + throw new Error('Too many malformed chunks in Ollama response'); + } + continue; + } + if (parsed.error) { + throw new Error(`Ollama error: ${parsed.error}`); + } + yield this.normalizeMessage(parsed.message); + } + } + if (buffer.trim() !== '') { + let parsed = null; + try { + parsed = this.parseChatResponse(buffer); + } + catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error); + utils_1.Logger.warn(`Skipped malformed chunk: ${buffer.slice(0, 50)}... - ${errorMessage}`, 'ollama-client'); + } + if (parsed?.error) { + throw new Error(`Ollama error: ${parsed.error}`); + } + if (parsed?.message) { + yield this.normalizeMessage(parsed.message); + } + } + } + catch (error) { + if (retryCount < this.maxRetries && this.isRetryableError(error, controller)) { + const errorMessage = error instanceof Error ? error.message : String(error); + utils_1.Logger.warn(`Retrying after error (attempt ${retryCount + 1}): ${errorMessage}`, 'ollama-client'); + await new Promise((resolve) => setTimeout(resolve, 10 * Math.pow(2, retryCount))); + yield* this.streamChatWithRetry(messages, tools, retryCount + 1); + } + else { + throw error; + } + } + finally { + reader?.releaseLock(); + if (this.currentStreamController === controller) { + this.currentStreamController = null; + } } - } - } catch (error) { - if (retryCount < this.maxRetries && !controller.signal.aborted) { - console.warn(`Retrying after error (attempt ${retryCount + 1}):`, error.message); - await new Promise((resolve) => setTimeout(resolve, 1000 * Math.pow(2, retryCount))); - yield* this.streamChatWithRetry(messages, tools, retryCount + 1); - } else { - throw error; - } - } finally { - this.currentStreamController = null; } - } - async chatWithRetry(messages, tools, retryCount) { - try { - const response = await this.fetchFn(`${this.baseURL}/api/chat`, { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - }, - body: JSON.stringify({ - model: this.model, - messages: messages, - tools: tools, - }), - }); - if (!response.ok) { - const text = await response.text(); - const parsed = JSON.parse(text); - const errorMsg = - typeof parsed.error === 'string' ? parsed.error : JSON.stringify(parsed.error); - throw new Error(`HTTP error! status: ${response.status}, message: ${errorMsg}`); - } - const data = await response.json(); - return data.message; - } catch (error) { - if (retryCount < this.maxRetries) { - console.warn(`Retrying after error (attempt ${retryCount + 1}):`, error.message); - await new Promise((resolve) => setTimeout(resolve, 1000 * Math.pow(2, retryCount))); - return this.chatWithRetry(messages, tools, retryCount + 1); - } else { - throw error; - } + async chatWithRetry(messages, tools = [], retryCount) { + const controller = new AbortController(); + this.currentStreamController = controller; + try { + const response = await this.fetchFn(`${this.baseURL}/api/chat`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + model: this.model, + messages: messages, + tools: tools, + stream: false, + }), + signal: controller.signal, + }); + if (!response.ok) { + throw new types_1.ApiError(`Ollama API error: ${response.status}`, response.status); + } + const data = await response.json(); + if (!this.isChatResponse(data)) { + return this.normalizeMessage(); + } + return this.normalizeMessage(data.message); + } + catch (error) { + if (retryCount < this.maxRetries && this.isRetryableError(error, controller)) { + const errorMessage = error instanceof Error ? error.message : String(error); + utils_1.Logger.warn(`Retrying after error (attempt ${retryCount + 1}): ${errorMessage}`, 'ollama-client'); + await new Promise((resolve) => setTimeout(resolve, 10 * Math.pow(2, retryCount))); + return this.chatWithRetry(messages, tools, retryCount + 1); + } + else { + throw error; + } + } + finally { + if (this.currentStreamController === controller) { + this.currentStreamController = null; + } + } + } + normalizeMessage(message) { + return { + role: message?.role ?? 'assistant', + content: message?.content ?? '', + tool_calls: message?.tool_calls ?? [], + tool_call_id: message?.tool_call_id, + }; + } + parseChatResponse(raw) { + const parsed = JSON.parse(raw); + if (!this.isChatResponse(parsed)) { + throw new Error('Invalid chat response'); + } + return parsed; + } + isChatResponse(data) { + if (typeof data !== 'object' || data === null) { + return false; + } + const response = data; + return ((response.error === undefined || typeof response.error === 'string') && + (response.message === undefined || this.isPartialMessage(response.message))); + } + isPartialMessage(data) { + if (typeof data !== 'object' || data === null) { + return false; + } + const message = data; + const validRole = message.role === undefined || + message.role === 'system' || + message.role === 'user' || + message.role === 'assistant' || + message.role === 'tool'; + return (validRole && + (message.content === undefined || typeof message.content === 'string') && + (message.tool_calls === undefined || Array.isArray(message.tool_calls)) && + (message.tool_call_id === undefined || typeof message.tool_call_id === 'string')); + } + isRetryableError(error, controller) { + if (controller.signal.aborted) { + return false; + } + if (error instanceof types_1.ApiError && error.statusCode >= 400 && error.statusCode < 500) { + return false; + } + if (error instanceof Error) { + if (error.name === 'AbortError') { + return false; + } + if (error.message.startsWith('Ollama error:') || + error.message.includes('Too many malformed chunks') || + error.message === 'No response body' || + error.message === 'Invalid response format') { + return false; + } + } + return true; } - } } exports.OllamaClient = OllamaClient; diff --git a/src/ollama-client.ts b/src/ollama-client.ts index 5d60322..b80b8d1 100644 --- a/src/ollama-client.ts +++ b/src/ollama-client.ts @@ -1,6 +1,6 @@ // src/ollama-client.ts -import type { OllamaMessage, OllamaTool } from './types'; +import type { CacheConfig, OllamaMessage, OllamaTool } from './types'; import { ApiError } from './types'; import { Logger } from './utils'; import { SemanticCacheService } from './semantic-cache'; @@ -15,16 +15,18 @@ export class OllamaClient { private model: string; private fetchFn: typeof fetch; private readonly maxRetries: number = 3; + private readonly maxMalformedChunks: number = 50; private currentStreamController: AbortController | null = null; private cacheService?: SemanticCacheService; - constructor(baseURL: string, model: string, fetchFn?: typeof fetch, cacheConfig?: any) { + constructor(baseURL: string, model: string, fetchFn?: typeof fetch, cacheConfig?: CacheConfig) { this.baseURL = baseURL; this.model = model; this.fetchFn = fetchFn ?? fetch; if (cacheConfig?.enabled) { this.cacheService = new SemanticCacheService(baseURL, cacheConfig); + void this.cacheService.initialize(); } } @@ -90,7 +92,7 @@ export class OllamaClient { if (lastUserMsg && this.cacheService) { const cached = await this.cacheService.getCache(lastUserMsg.content); if (cached) { - return { content: cached }; + return { role: 'assistant', content: cached }; } } @@ -101,6 +103,25 @@ export class OllamaClient { return response; } + async streamChatAsPromise( + messages: OllamaMessage[], + tools: OllamaTool[] = [] + ): Promise { + let content = ''; + let role: OllamaMessage['role'] = 'assistant'; + let toolCalls: OllamaMessage['tool_calls']; + + for await (const chunk of this.streamChat(messages, tools)) { + role = chunk.role ?? role; + content += chunk.content ?? ''; + if (chunk.tool_calls) { + toolCalls = [...(toolCalls ?? []), ...chunk.tool_calls]; + } + } + + return { role, content, tool_calls: toolCalls }; + } + async *streamChatWithRetry( messages: OllamaMessage[], tools: OllamaTool[] = [], @@ -108,6 +129,7 @@ export class OllamaClient { ): AsyncGenerator { const controller = new AbortController(); this.currentStreamController = controller; + let reader: ReadableStreamDefaultReader | null = null; try { const response = await this.fetchFn(`${this.baseURL}/api/chat`, { @@ -118,23 +140,29 @@ export class OllamaClient { body: JSON.stringify({ model: this.model, messages: messages, - stream: true, tools: tools, + stream: true, }), signal: controller.signal, }); if (!response.ok) { - throw new Error(`HTTP error! status: ${response.status}`); + throw new ApiError(`Ollama API error: ${response.status}`, response.status); } if (!response.body) { - throw new Error('Response body is null'); + throw new Error('No response body'); } - const reader = response.body.getReader(); + const contentType = response.headers?.get?.('content-type'); + if (contentType && !contentType.includes('application/x-ndjson')) { + throw new Error('Invalid response format'); + } + + reader = response.body.getReader(); const decoder = new TextDecoder(); let buffer = ''; + let malformedChunks = 0; while (true) { const { done, value } = await reader.read(); @@ -151,34 +179,64 @@ export class OllamaClient { continue; } + let parsed: OllamaChatResponse; try { - const parsed = JSON.parse(line); - yield parsed.message; + parsed = this.parseChatResponse(line); } catch (error) { - // Log but don't throw - malformed JSON is not critical - console.error('Failed to parse chunk:', line, error); + malformedChunks++; + const errorMessage = error instanceof Error ? error.message : String(error); + Logger.warn( + `Skipped malformed chunk: ${line.slice(0, 50)}... - ${errorMessage}`, + 'ollama-client' + ); + if (malformedChunks > this.maxMalformedChunks) { + throw new Error('Too many malformed chunks in Ollama response'); + } + continue; } + + if (parsed.error) { + throw new Error(`Ollama error: ${parsed.error}`); + } + + yield this.normalizeMessage(parsed.message); } } if (buffer.trim() !== '') { + let parsed: OllamaChatResponse | null = null; try { - const parsed = JSON.parse(buffer); - yield parsed.message; + parsed = this.parseChatResponse(buffer); } catch (error) { - console.error('Failed to parse final chunk:', buffer, error); + const errorMessage = error instanceof Error ? error.message : String(error); + Logger.warn( + `Skipped malformed chunk: ${buffer.slice(0, 50)}... - ${errorMessage}`, + 'ollama-client' + ); + } + + if (parsed?.error) { + throw new Error(`Ollama error: ${parsed.error}`); + } + + if (parsed?.message) { + yield this.normalizeMessage(parsed.message); } } } catch (error) { - if (retryCount < this.maxRetries && !controller.signal.aborted) { - console.warn(`Retrying after error (attempt ${retryCount + 1}):`, error.message); - await new Promise((resolve) => setTimeout(resolve, 1000 * Math.pow(2, retryCount))); + if (retryCount < this.maxRetries && this.isRetryableError(error, controller)) { + const errorMessage = error instanceof Error ? error.message : String(error); + Logger.warn(`Retrying after error (attempt ${retryCount + 1}): ${errorMessage}`, 'ollama-client'); + await new Promise((resolve) => setTimeout(resolve, 10 * Math.pow(2, retryCount))); yield* this.streamChatWithRetry(messages, tools, retryCount + 1); } else { throw error; } } finally { - this.currentStreamController = null; + reader?.releaseLock(); + if (this.currentStreamController === controller) { + this.currentStreamController = null; + } } } @@ -187,6 +245,9 @@ export class OllamaClient { tools: OllamaTool[] = [], retryCount: number ): Promise { + const controller = new AbortController(); + this.currentStreamController = controller; + try { const response = await this.fetchFn(`${this.baseURL}/api/chat`, { method: 'POST', @@ -197,27 +258,115 @@ export class OllamaClient { model: this.model, messages: messages, tools: tools, + stream: false, }), + signal: controller.signal, }); if (!response.ok) { - const text = await response.text(); - const parsed = JSON.parse(text); - const errorMsg = - typeof parsed.error === 'string' ? parsed.error : JSON.stringify(parsed.error); - throw new Error(`HTTP error! status: ${response.status}, message: ${errorMsg}`); + throw new ApiError(`Ollama API error: ${response.status}`, response.status); } - const data = await response.json(); - return data.message; + const data = await response.json() as unknown; + if (!this.isChatResponse(data)) { + return this.normalizeMessage(); + } + return this.normalizeMessage(data.message); } catch (error) { - if (retryCount < this.maxRetries) { - console.warn(`Retrying after error (attempt ${retryCount + 1}):`, error.message); - await new Promise((resolve) => setTimeout(resolve, 1000 * Math.pow(2, retryCount))); + if (retryCount < this.maxRetries && this.isRetryableError(error, controller)) { + const errorMessage = error instanceof Error ? error.message : String(error); + Logger.warn(`Retrying after error (attempt ${retryCount + 1}): ${errorMessage}`, 'ollama-client'); + await new Promise((resolve) => setTimeout(resolve, 10 * Math.pow(2, retryCount))); return this.chatWithRetry(messages, tools, retryCount + 1); } else { throw error; } + } finally { + if (this.currentStreamController === controller) { + this.currentStreamController = null; + } } } + + private normalizeMessage(message?: Partial): OllamaMessage { + return { + role: message?.role ?? 'assistant', + content: message?.content ?? '', + tool_calls: message?.tool_calls ?? [], + tool_call_id: message?.tool_call_id, + }; + } + + private parseChatResponse(raw: string): OllamaChatResponse { + const parsed: unknown = JSON.parse(raw); + if (!this.isChatResponse(parsed)) { + throw new Error('Invalid chat response'); + } + return parsed; + } + + private isChatResponse(data: unknown): data is OllamaChatResponse { + if (typeof data !== 'object' || data === null) { + return false; + } + + const response = data as { message?: unknown; error?: unknown }; + return ( + (response.error === undefined || typeof response.error === 'string') && + (response.message === undefined || this.isPartialMessage(response.message)) + ); + } + + private isPartialMessage(data: unknown): data is Partial { + if (typeof data !== 'object' || data === null) { + return false; + } + + const message = data as { + role?: unknown; + content?: unknown; + tool_calls?: unknown; + tool_call_id?: unknown; + }; + const validRole = + message.role === undefined || + message.role === 'system' || + message.role === 'user' || + message.role === 'assistant' || + message.role === 'tool'; + + return ( + validRole && + (message.content === undefined || typeof message.content === 'string') && + (message.tool_calls === undefined || Array.isArray(message.tool_calls)) && + (message.tool_call_id === undefined || typeof message.tool_call_id === 'string') + ); + } + + private isRetryableError(error: unknown, controller: AbortController): boolean { + if (controller.signal.aborted) { + return false; + } + + if (error instanceof ApiError && error.statusCode >= 400 && error.statusCode < 500) { + return false; + } + + if (error instanceof Error) { + if (error.name === 'AbortError') { + return false; + } + + if ( + error.message.startsWith('Ollama error:') || + error.message.includes('Too many malformed chunks') || + error.message === 'No response body' || + error.message === 'Invalid response format' + ) { + return false; + } + } + + return true; + } } diff --git a/src/semantic-cache.js b/src/semantic-cache.js index 17187f8..9b12c59 100644 --- a/src/semantic-cache.js +++ b/src/semantic-cache.js @@ -1,96 +1,105 @@ -'use strict'; +"use strict"; +// eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-unsafe-call, @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-unsafe-return +/* eslint-disable */ // src/semantic-cache.ts -Object.defineProperty(exports, '__esModule', { value: true }); +Object.defineProperty(exports, "__esModule", { value: true }); exports.SemanticCacheService = void 0; -const chromadb_1 = require('chromadb'); -const utils_1 = require('./utils'); +const chromadb_1 = require("chromadb"); +const utils_1 = require("./utils"); class SemanticCacheService { - constructor(ollamaURL, config) { - this.collection = null; - this.ollamaURL = ollamaURL.replace(/\/+$/, ''); - this.config = config; - // Use configurable ChromaDB URL or default to localhost - this.chromaURL = config.chromaURL || 'http://localhost:8000'; - this.client = new chromadb_1.ChromaClient({ path: this.chromaURL }); - } - async initialize() { - if (!this.config.enabled) return; - try { - this.collection = await this.client.getOrCreateCollection({ - name: this.config.collectionName, - metadata: { 'hnsw:space': 'cosine' }, - }); - utils_1.Logger.info( - `Semantic cache initialized: ${this.config.collectionName}`, - 'semantic-cache' - ); - } catch (error) { - utils_1.Logger.error( - `Failed to initialize semantic cache: ${error.message}`, - 'semantic-cache' - ); - throw error; + constructor(ollamaURL, config) { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + this.collection = null; + this.ollamaURL = ollamaURL.replace(/\/+$/, ''); + this.config = config; + // Use configurable ChromaDB URL or default to localhost + const chromaURL = config.chromaURL || 'http://localhost:8000'; + this.client = new chromadb_1.ChromaClient({ path: chromaURL }); } - } - async getCache(query) { - if (!this.config.enabled || !this.collection) return null; - try { - const results = await this.collection.query({ - query_embeddings: await this.generateEmbedding(query), - n_results: 1, - where: { source: 'ollama' }, - }); - if (results.ids[0] && results.ids[0].length > 0) { - const [id] = results.ids[0]; - const [content] = results.documents[0]; - if (results.distances[0] && results.distances[0][0] > this.config.similarityThreshold) { - return content; + async initialize() { + if (!this.config.enabled) + return; + try { + this.collection = await this.client.getOrCreateCollection({ + name: this.config.collectionName, + metadata: { 'hnsw:space': 'cosine' }, + }); + utils_1.Logger.info(`Semantic cache initialized: ${this.config.collectionName}`, 'semantic-cache'); + } + catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error); + utils_1.Logger.error(`Failed to initialize semantic cache: ${errorMessage}`, 'semantic-cache'); + throw error; } - } - return null; - } catch (error) { - utils_1.Logger.warn(`Cache lookup failed: ${error.message}`, 'semantic-cache'); - return null; } - } - async setCache(query, response) { - if (!this.config.enabled || !this.collection) return; - try { - await this.collection.add({ - ids: [crypto.randomUUID()], - documents: [response], - embeddings: await this.generateEmbedding(query), - metadatas: [{ source: 'ollama' }], - }); - } catch (error) { - utils_1.Logger.warn(`Cache set failed: ${error.message}`, 'semantic-cache'); + async getCache(query) { + if (!this.config.enabled || !this.collection) + return null; + try { + const results = await this.collection.query({ + query_embeddings: await this.generateEmbedding(query), + n_results: 1, + where: { source: 'ollama' }, + }); + if (results.ids[0] && results.ids[0].length > 0) { + const [id] = results.ids[0]; + const [content] = results.documents[0]; + if (results.distances[0] && results.distances[0][0] > this.config.similarityThreshold) { + return content; + } + } + return null; + } + catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error); + utils_1.Logger.warn(`Cache lookup failed: ${errorMessage}`, 'semantic-cache'); + return null; + } } - } - async clearCache() { - if (!this.config.enabled || !this.collection) return; - try { - await this.collection.reset(); - utils_1.Logger.info('Semantic cache cleared', 'semantic-cache'); - } catch (error) { - utils_1.Logger.error(`Failed to clear semantic cache: ${error.message}`, 'semantic-cache'); + async setCache(query, response) { + if (!this.config.enabled || !this.collection) + return; + try { + await this.collection.add({ + ids: [crypto.randomUUID()], + documents: [response], + embeddings: await this.generateEmbedding(query), + metadatas: [{ source: 'ollama' }], + }); + } + catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error); + utils_1.Logger.warn(`Cache set failed: ${errorMessage}`, 'semantic-cache'); + } } - } - async generateEmbedding(text) { - const response = await fetch(`${this.ollamaURL}/api/embeddings`, { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - }, - body: JSON.stringify({ - model: this.config.embeddingModel, - prompt: text, - }), - }); - if (!response.ok) { - throw new Error(`Failed to generate embedding: ${response.status} ${response.statusText}`); + async clearCache() { + if (!this.config.enabled || !this.collection) + return; + try { + await this.collection.reset(); + utils_1.Logger.info('Semantic cache cleared', 'semantic-cache'); + } + catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error); + utils_1.Logger.error(`Failed to clear semantic cache: ${errorMessage}`, 'semantic-cache'); + } + } + async generateEmbedding(text) { + const response = await fetch(`${this.ollamaURL}/api/embeddings`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + model: this.config.embeddingModel, + prompt: text, + }), + }); + if (!response.ok) { + throw new Error(`Failed to generate embedding: ${response.status} ${response.statusText}`); + } + const data = await response.json(); + return data.embedding; } - const data = await response.json(); - return data.embedding; - } } exports.SemanticCacheService = SemanticCacheService; diff --git a/src/tool-executor.js b/src/tool-executor.js index b80f258..89ec8d3 100644 --- a/src/tool-executor.js +++ b/src/tool-executor.js @@ -2,6 +2,7 @@ // src/tool-executor.ts Object.defineProperty(exports, "__esModule", { value: true }); exports.ToolExecutor = void 0; +const obsidian_1 = require("obsidian"); const utils_1 = require("./utils"); // Disallow characters that are invalid in file paths const INVALID_PATH_CHARS = /[<>:"|?*~]/; @@ -80,6 +81,10 @@ class ToolExecutor { switch (toolName) { case 'create_file': return await this.handleCreateFile(parsedArgs); + case 'read_vault_file': + return await this.handleReadVaultFile(parsedArgs); + case 'search_vault_files': + return this.handleSearchVaultFiles(parsedArgs); default: return { success: false, message: `Unknown tool: ${toolName}` }; } @@ -110,5 +115,53 @@ class ToolExecutor { throw new Error(errorMessage); } } + async executeTool(name, args) { + return this.handleToolCall({ + id: crypto.randomUUID(), + type: 'function', + function: { + name, + arguments: args, + }, + }); + } + async handleReadVaultFile(args) { + const path = args.path; + if (typeof path !== 'string') { + throw new Error('Path must be a string'); + } + if (!this.isSafePath(path)) { + throw new Error('Invalid file path detected'); + } + const file = this.vault.getAbstractFileByPath(path); + if (!(file instanceof obsidian_1.TFile)) { + throw new Error(`File not found: ${path}`); + } + const content = await this.vault.cachedRead(file); + return { + success: true, + message: 'File read successfully', + data: { path, content }, + }; + } + handleSearchVaultFiles(args) { + const query = args.query; + const limitArg = args.limit; + if (typeof query !== 'string') { + throw new Error('Query must be a string'); + } + const limit = typeof limitArg === 'number' && Number.isFinite(limitArg) ? limitArg : 10; + const normalizedQuery = query.toLowerCase(); + const files = this.vault + .getMarkdownFiles() + .filter((file) => file.path.toLowerCase().includes(normalizedQuery)) + .slice(0, limit) + .map((file) => ({ path: file.path, basename: file.basename })); + return { + success: true, + message: `Found ${files.length} matching files`, + data: files, + }; + } } exports.ToolExecutor = ToolExecutor; diff --git a/src/tool-executor.ts b/src/tool-executor.ts index c81bb11..5b64ff8 100644 --- a/src/tool-executor.ts +++ b/src/tool-executor.ts @@ -1,6 +1,6 @@ // src/tool-executor.ts -import { Vault, App } from 'obsidian'; +import { Vault, App, TFile } from 'obsidian'; import type { ToolCall, ToolResult } from './types'; import { safeParseJson } from './utils'; @@ -95,6 +95,10 @@ export class ToolExecutor { switch (toolName) { case 'create_file': return await this.handleCreateFile(parsedArgs); + case 'read_vault_file': + return await this.handleReadVaultFile(parsedArgs); + case 'search_vault_files': + return this.handleSearchVaultFiles(parsedArgs); default: return { success: false, message: `Unknown tool: ${toolName}` }; } @@ -128,4 +132,62 @@ export class ToolExecutor { throw new Error(errorMessage); } } + + async executeTool(name: string, args: string | Record): Promise { + return this.handleToolCall({ + id: crypto.randomUUID(), + type: 'function', + function: { + name, + arguments: args as string, + }, + }); + } + + private async handleReadVaultFile(args: Record): Promise { + const path = args.path; + + if (typeof path !== 'string') { + throw new Error('Path must be a string'); + } + + if (!this.isSafePath(path)) { + throw new Error('Invalid file path detected'); + } + + const file = this.vault.getAbstractFileByPath(path); + if (!(file instanceof TFile)) { + throw new Error(`File not found: ${path}`); + } + + const content = await this.vault.cachedRead(file); + return { + success: true, + message: 'File read successfully', + data: { path, content }, + }; + } + + private handleSearchVaultFiles(args: Record): ToolResult { + const query = args.query; + const limitArg = args.limit; + + if (typeof query !== 'string') { + throw new Error('Query must be a string'); + } + + const limit = typeof limitArg === 'number' && Number.isFinite(limitArg) ? limitArg : 10; + const normalizedQuery = query.toLowerCase(); + const files = this.vault + .getMarkdownFiles() + .filter((file) => file.path.toLowerCase().includes(normalizedQuery)) + .slice(0, limit) + .map((file) => ({ path: file.path, basename: file.basename })); + + return { + success: true, + message: `Found ${files.length} matching files`, + data: files, + }; + } } diff --git a/src/types.js b/src/types.js index 9bd89f2..2e96d54 100644 --- a/src/types.js +++ b/src/types.js @@ -55,7 +55,7 @@ class StreamingError extends OllamaError { } exports.StreamingError = StreamingError; class ToolExecutionError extends OllamaError { - constructor(message, toolName) { + constructor(message, toolName = 'unknown') { super(message, ErrorType.TOOL_EXECUTION_ERROR); this.toolName = toolName; Object.setPrototypeOf(this, ToolExecutionError.prototype); @@ -63,7 +63,7 @@ class ToolExecutionError extends OllamaError { } exports.ToolExecutionError = ToolExecutionError; class PathValidationError extends OllamaError { - constructor(message, path) { + constructor(message, path = '') { super(message, ErrorType.PATH_VALIDATION_ERROR); this.path = path; Object.setPrototypeOf(this, PathValidationError.prototype); diff --git a/src/types.ts b/src/types.ts index d1bd0e3..4f76f11 100644 --- a/src/types.ts +++ b/src/types.ts @@ -30,6 +30,7 @@ export class NetworkError extends OllamaError { constructor(message: string, statusCode?: number) { super(message, ErrorType.NETWORK_ERROR); this.statusCode = statusCode; + Object.setPrototypeOf(this, NetworkError.prototype); } } @@ -39,24 +40,44 @@ export class ApiError extends OllamaError { constructor(message: string, statusCode: number) { super(message, ErrorType.API_ERROR); this.statusCode = statusCode; + Object.setPrototypeOf(this, ApiError.prototype); } } export class ValidationError extends OllamaError { - constructor(message: string) { + public readonly details?: { field?: string; message?: string }; + + constructor(message: string, details?: { field?: string; message?: string }) { super(message, ErrorType.VALIDATION_ERROR); + this.details = details; + Object.setPrototypeOf(this, ValidationError.prototype); } } export class StreamingError extends OllamaError { constructor(message: string) { super(message, ErrorType.STREAMING_ERROR); + Object.setPrototypeOf(this, StreamingError.prototype); } } export class ToolExecutionError extends OllamaError { - constructor(message: string) { + public readonly toolName: string; + + constructor(message: string, toolName: string = 'unknown') { super(message, ErrorType.TOOL_EXECUTION_ERROR); + this.toolName = toolName; + Object.setPrototypeOf(this, ToolExecutionError.prototype); + } +} + +export class PathValidationError extends OllamaError { + public readonly path: string; + + constructor(message: string, path: string = '') { + super(message, ErrorType.PATH_VALIDATION_ERROR); + this.path = path; + Object.setPrototypeOf(this, PathValidationError.prototype); } } @@ -65,9 +86,10 @@ export class ToolExecutionError extends OllamaError { // ============================================================ export interface OllamaMessage { - role: 'system' | 'user' | 'assistant'; + role: 'system' | 'user' | 'assistant' | 'tool'; content: string; tool_calls?: OllamaToolCall[]; + tool_call_id?: string; } export interface OllamaToolCall { @@ -97,6 +119,31 @@ export interface OllamaTool { }; } +export type ToolCall = OllamaToolCall; + +export interface ToolResult { + success: boolean; + message: string; + data?: unknown; + id?: string; +} + +export interface VaultIndexEntry { + path: string; + title: string; + content: string; + score: number; +} + +export interface ChatMessage { + id: string; + role: 'user' | 'assistant'; + content: string; + timestamp: number; + isStreaming?: boolean; + tool_calls?: OllamaToolCall[]; +} + // ============================================================ // Plugin Configuration // ============================================================ diff --git a/src/utils.js b/src/utils.js index 889e40e..69ff631 100644 --- a/src/utils.js +++ b/src/utils.js @@ -148,8 +148,9 @@ function safeParseJson(jsonString) { return true; } // Recursively check nested objects (own properties only) + const record = obj; for (const key of Object.keys(obj)) { - if (checkDangerousPatterns(obj[key])) { + if (checkDangerousPatterns(record[key])) { return true; } } diff --git a/src/utils.ts b/src/utils.ts index 75d4cba..64a420b 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -167,8 +167,9 @@ export function safeParseJson(jsonString: string): unknown { } // Recursively check nested objects (own properties only) - for (const key of Object.keys(obj as Record)) { - if (checkDangerousPatterns((obj as Record)[key])) { + const record = obj as Record; + for (const key of Object.keys(obj)) { + if (checkDangerousPatterns(record[key])) { return true; } } diff --git a/src/vault-indexer.js b/src/vault-indexer.js index 103fb8d..ea9b842 100644 --- a/src/vault-indexer.js +++ b/src/vault-indexer.js @@ -1,151 +1,226 @@ -'use strict'; +"use strict"; // src/vault-indexer.ts -Object.defineProperty(exports, '__esModule', { value: true }); -exports.VaultIndexer = void 0; -const obsidian_1 = require('obsidian'); -const utils_1 = require('./utils'); +Object.defineProperty(exports, "__esModule", { value: true }); +exports.VaultIndexer = exports.InMemoryCache = void 0; +const utils_1 = require("./utils"); +class InMemoryCache { + constructor() { + this.store = new Map(); + } + get(key) { + return Promise.resolve(this.store.get(key) ?? null); + } + put(key, value) { + this.store.set(key, value); + return Promise.resolve(); + } +} +exports.InMemoryCache = InMemoryCache; +const STOP_WORDS = new Set([ + 'a', 'an', 'the', 'is', 'it', 'in', 'on', 'at', 'to', 'for', 'of', + 'and', 'or', 'but', 'with', 'by', 'from', 'up', 'about', 'into', + 'this', 'that', 'these', 'those', 'be', 'been', 'being', 'have', + 'has', 'had', 'do', 'does', 'did', 'will', 'would', 'could', 'should', + 'may', 'might', 'can', 'are', 'was', 'were', 'as', 'so', 'if', 'not', + 'no', 'my', 'your', 'our', 'its', 'we', 'you', 'he', 'she', 'they', +]); +const CONTENT_PREVIEW_LENGTH = 500; class VaultIndexer { - constructor(vault) { - this.vault = vault; - this.stemmer = new obsidian_1.PorterStemmer(); - this.SCORING_WEIGHTS = { - TITLE: 5, - FRONTMATTER_TITLE: 4, - FRONTMATTER_TAGS: 3, - HEADINGS: 2, - CONTENT: 1, - }; - } - async getVaultEntries() { - const files = this.vault.getMarkdownFiles(); - const entries = []; - for (const file of files) { - try { - const content = await this.vault.cachedRead(file); + constructor(vault, cache) { + this.SCORING_WEIGHTS = { + TITLE: 5, + FRONTMATTER_TITLE: 4, + FRONTMATTER_TAGS: 3, + HEADINGS: 2, + CONTENT: 1, + }; + this.vault = vault; + this.cache = cache; + } + tokenize(text) { + return text + .toLowerCase() + .replace(/[^\w\s]/g, '') + .split(/\s+/) + .filter((token) => token.length > 1 && !STOP_WORDS.has(token)); + } + tokenizeContent(content, file) { const parsed = this.parseMarkdown(content); - entries.push({ - file: file, - title: parsed.title, - frontmatter: parsed.frontmatter, - headings: parsed.headings, - content: parsed.content, - basename: file.basename, - }); - } catch (error) { - utils_1.Logger.warn(`Failed to read file ${file.path}: ${error.message}`, 'vault-indexer'); - } + const bodyWithoutFrontmatter = content.replace(/^---[\s\S]*?---\n/, ''); + const paragraphs = bodyWithoutFrontmatter + .split(/\n\n+/) + .map((p) => p.trim()) + .filter((p) => p && !p.startsWith('#')); + const firstParagraph = paragraphs[0] || ''; + return { + title: parsed.title || file.basename, + headings: parsed.headings, + frontmatter: parsed.frontmatter, + firstParagraph, + content: parsed.content, + basename: file.basename, + }; } - return entries; - } - async searchVault(query, limit = 3) { - const entries = await this.getVaultEntries(); - const tokenized = this.tokenizeQuery(query); - const scoredEntries = entries.map((entry) => { - let score = 0; - for (const queryToken of tokenized.tokens) { - const stemmed = this.stemToken(queryToken); - let matched = false; - if (entry.frontmatter?.title && this.exactMatch(entry.frontmatter.title, queryToken)) { - score += this.SCORING_WEIGHTS.FRONTMATTER_TITLE; - matched = true; - } else if ( - entry.basename && - this.exactMatch(entry.basename.replace(/\.md$/, ''), queryToken) - ) { - score += this.SCORING_WEIGHTS.FRONTMATTER_TITLE; - matched = true; - } - if (entry.frontmatter?.tags && this.exactMatch(entry.frontmatter.tags, queryToken)) { - score += this.SCORING_WEIGHTS.FRONTMATTER_TAGS; - matched = true; - } - if (entry.headings.some((heading) => heading.toLowerCase().includes(stemmed))) { - score += this.SCORING_WEIGHTS.HEADINGS; - matched = true; - } - if (entry.content.toLowerCase().includes(stemmed)) { - score += this.SCORING_WEIGHTS.CONTENT; - matched = true; - } - if (matched) { - // Add bonus for exact matches - if (entry.title && this.exactMatch(entry.title, queryToken)) { - score += this.SCORING_WEIGHTS.TITLE; - } - } - } - return { entry, score }; - }); - // Sort by score descending and return top results - scoredEntries.sort((a, b) => b.score - a.score); - return scoredEntries.slice(0, limit).map((item) => item.entry); - } - tokenizeQuery(query) { - const tokens = query - .toLowerCase() - .replace(/[^\w\s]/g, '') - .split(/\s+/) - .filter((token) => token.length > 0); - return { - tokens, - stemmedTokens: tokens.map((token) => this.stemToken(token)), - }; - } - stemToken(token) { - return this.stemmer.stemWord(token); - } - exactMatch(text, queryToken) { - if (!text) return false; - return text.toLowerCase().includes(queryToken.toLowerCase()); - } - parseMarkdown(content) { - const frontmatterRegex = /^---\n([\s\S]*?)\n---\n/; - const frontmatterMatch = content.match(frontmatterRegex); - const frontmatter = {}; - if (frontmatterMatch) { - try { - const frontmatterContent = frontmatterMatch[1]; - const lines = frontmatterContent.trim().split('\n'); - for (const line of lines) { - const [key, ...valueParts] = line.split(':'); - if (!key) continue; - const value = valueParts.join(':').trim(); - if (key.trim() === 'title') { - if (value) { - frontmatter.title = value; + calculateWeightedScore(tokenized, queryTokens) { + let score = 0; + for (const token of queryTokens) { + if (tokenized.frontmatter?.title && this.exactMatch(tokenized.frontmatter.title, token)) { + score += this.SCORING_WEIGHTS.FRONTMATTER_TITLE; } - } else if (key.trim() === 'tags') { - if (value) { - frontmatter.tags = value; + if (tokenized.basename && this.exactMatch(tokenized.basename, token)) { + score += this.SCORING_WEIGHTS.FRONTMATTER_TITLE; + } + if (tokenized.frontmatter?.tags && this.exactMatch(tokenized.frontmatter.tags, token)) { + score += this.SCORING_WEIGHTS.FRONTMATTER_TAGS; + } + if (tokenized.headings.some((h) => h.toLowerCase().includes(token.toLowerCase()))) { + score += this.SCORING_WEIGHTS.HEADINGS; + } + if (tokenized.content.toLowerCase().includes(token.toLowerCase())) { + score += this.SCORING_WEIGHTS.CONTENT; + } + if (tokenized.title && this.exactMatch(tokenized.title, token)) { + score += this.SCORING_WEIGHTS.TITLE; } - } } - } catch { - utils_1.Logger.warn('Failed to parse frontmatter', 'vault-indexer'); - } + return { score }; } - const titleMatch = content.match(/^# (.+)$/m); - const title = titleMatch ? titleMatch[1] : ''; - const headings = []; - const headingRegex = /^#{1,6} (.+)$/gm; - let headingMatch; - while ((headingMatch = headingRegex.exec(content)) !== null) { - headings.push(headingMatch[1]); + async getVaultEntries() { + const files = this.vault.getMarkdownFiles(); + const entries = []; + for (const file of files) { + try { + const content = typeof this.vault.cachedRead === 'function' + ? await this.vault.cachedRead(file) + : await this.vault.read(file); + const parsed = this.parseMarkdown(content); + entries.push({ + file: file, + title: parsed.frontmatter.title || file.basename, + frontmatter: parsed.frontmatter, + headings: parsed.headings, + content: parsed.content.slice(0, CONTENT_PREVIEW_LENGTH), + basename: file.basename, + score: 0, + }); + } + catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error); + utils_1.Logger.warn(`Failed to read file ${file.path}: ${errorMessage}`, 'vault-indexer'); + } + } + return entries; + } + async searchVault(query, limit = 3) { + if (!query || !query.trim()) { + return []; + } + const cacheKey = `query:${query.trim()}:limit:${limit}`; + if (this.cache) { + let cachedResults = null; + try { + cachedResults = await this.cache.get(cacheKey); + } + catch { + cachedResults = null; + } + if (cachedResults) { + try { + // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment + const parsedResults = JSON.parse(cachedResults); + return parsedResults.slice(0, limit); + } + catch { + // ignore parse errors + } + } + } + const queryTokens = this.tokenize(query); + if (queryTokens.length === 0) { + return []; + } + const entries = await this.getVaultEntries(); + const scored = entries + .map((entry) => { + const { score } = this.calculateWeightedScore({ + title: entry.title, + headings: entry.headings, + frontmatter: entry.frontmatter, + firstParagraph: '', + content: entry.content, + basename: entry.basename, + }, queryTokens); + return { ...entry, score }; + }) + .filter((e) => e.score > 0); + scored.sort((a, b) => b.score - a.score); + const results = scored.slice(0, limit); + if (this.cache) { + try { + await this.cache.put(cacheKey, JSON.stringify(results)); + } + catch (error) { + const errorMessage = error instanceof Error ? error.message : String(error); + utils_1.Logger.warn(`Failed to cache results for query "${query}": ${errorMessage}`, 'vault-indexer'); + } + } + return results; + } + stemToken(token) { + if (token.endsWith('ing') && token.length > 4) + return token.slice(0, -3); + if (token.endsWith('ed') && token.length > 3) + return token.slice(0, -2); + if (token.endsWith('s') && token.length > 2) + return token.slice(0, -1); + return token; + } + exactMatch(text, queryToken) { + if (!text) + return false; + const textLower = text.toLowerCase(); + const queryLower = queryToken.toLowerCase(); + const queryStem = this.stemToken(queryLower); + return textLower.includes(queryLower) || textLower.includes(queryStem); + } + parseMarkdown(content) { + const frontmatterRegex = /^---\n([\s\S]*?)\n---\n/; + const frontmatterMatch = content.match(frontmatterRegex); + const frontmatter = {}; + if (frontmatterMatch) { + try { + const lines = frontmatterMatch[1].trim().split('\n'); + for (const line of lines) { + const [key, ...valueParts] = line.split(':'); + if (!key) + continue; + const value = valueParts.join(':').trim(); + if (key.trim() === 'title' && value) + frontmatter.title = value; + else if (key.trim() === 'tags' && value) + frontmatter.tags = value; + } + } + catch { + utils_1.Logger.warn('Failed to parse frontmatter', 'vault-indexer'); + } + } + const titleMatch = content.match(/^# (.+)$/m); + const title = titleMatch ? titleMatch[1] : ''; + const headings = []; + const headingRegex = /^#{1,6} (.+)$/gm; + let headingMatch; + while ((headingMatch = headingRegex.exec(content)) !== null) { + headings.push(headingMatch[1]); + } + const bodyWithoutFrontmatter = frontmatterMatch + ? content.substring(frontmatterMatch[0].length) + : content; + const bodyText = bodyWithoutFrontmatter + .replace(/#{1,6} .+/g, '') + .replace(/^\s*[\r\n]/gm, '') + .trim(); + return { frontmatter, title, headings, content: bodyText }; } - const contentWithoutFrontmatter = frontmatterMatch - ? content.substring(frontmatterMatch[0].length) - : content; - const contentWithoutFrontmatterAndHeadings = contentWithoutFrontmatter - .replace(/#{1,6} .+/g, '') // Remove headings - .replace(/^---[\s\S]*?---\n/, '') // Remove frontmatter - .replace(/^\s*[\r\n]/gm, '') // Remove empty lines - .trim(); - return { - frontmatter, - title, - headings, - content: contentWithoutFrontmatterAndHeadings, - }; - } } exports.VaultIndexer = VaultIndexer; diff --git a/src/vault-indexer.ts b/src/vault-indexer.ts index 10daef8..4771fab 100644 --- a/src/vault-indexer.ts +++ b/src/vault-indexer.ts @@ -10,11 +10,13 @@ interface ParsedFrontmatter { [key: string]: unknown; } -interface ParsedMarkdown { - frontmatter: ParsedFrontmatter; +interface TokenizedContent { title: string; headings: string[]; + frontmatter: ParsedFrontmatter; + firstParagraph: string; content: string; + basename: string; } interface VaultEntry { @@ -24,8 +26,33 @@ interface VaultEntry { headings: string[]; content: string; basename: string; + score: number; } +export class InMemoryCache implements Cache { + private store = new Map(); + + get(key: string): Promise { + return Promise.resolve(this.store.get(key) ?? null); + } + + put(key: string, value: string): Promise { + this.store.set(key, value); + return Promise.resolve(); + } +} + +const STOP_WORDS = new Set([ + 'a', 'an', 'the', 'is', 'it', 'in', 'on', 'at', 'to', 'for', 'of', + 'and', 'or', 'but', 'with', 'by', 'from', 'up', 'about', 'into', + 'this', 'that', 'these', 'those', 'be', 'been', 'being', 'have', + 'has', 'had', 'do', 'does', 'did', 'will', 'would', 'could', 'should', + 'may', 'might', 'can', 'are', 'was', 'were', 'as', 'so', 'if', 'not', + 'no', 'my', 'your', 'our', 'its', 'we', 'you', 'he', 'she', 'they', +]); + +const CONTENT_PREVIEW_LENGTH = 500; + export class VaultIndexer { private vault: Vault; private cache?: Cache; @@ -42,20 +69,78 @@ export class VaultIndexer { this.cache = cache; } + tokenize(text: string): string[] { + return text + .toLowerCase() + .replace(/[^\w\s]/g, '') + .split(/\s+/) + .filter((token) => token.length > 1 && !STOP_WORDS.has(token)); + } + + tokenizeContent(content: string, file: TFile): TokenizedContent { + const parsed = this.parseMarkdown(content); + const bodyWithoutFrontmatter = content.replace(/^---[\s\S]*?---\n/, ''); + const paragraphs = bodyWithoutFrontmatter + .split(/\n\n+/) + .map((p) => p.trim()) + .filter((p) => p && !p.startsWith('#')); + const firstParagraph = paragraphs[0] || ''; + return { + title: parsed.title || file.basename, + headings: parsed.headings, + frontmatter: parsed.frontmatter, + firstParagraph, + content: parsed.content, + basename: file.basename, + }; + } + + calculateWeightedScore( + tokenized: TokenizedContent, + queryTokens: string[] + ): { score: number } { + let score = 0; + for (const token of queryTokens) { + if (tokenized.frontmatter?.title && this.exactMatch(tokenized.frontmatter.title, token)) { + score += this.SCORING_WEIGHTS.FRONTMATTER_TITLE; + } + if (tokenized.basename && this.exactMatch(tokenized.basename, token)) { + score += this.SCORING_WEIGHTS.FRONTMATTER_TITLE; + } + if (tokenized.frontmatter?.tags && this.exactMatch(tokenized.frontmatter.tags, token)) { + score += this.SCORING_WEIGHTS.FRONTMATTER_TAGS; + } + if (tokenized.headings.some((h) => h.toLowerCase().includes(token.toLowerCase()))) { + score += this.SCORING_WEIGHTS.HEADINGS; + } + if (tokenized.content.toLowerCase().includes(token.toLowerCase())) { + score += this.SCORING_WEIGHTS.CONTENT; + } + if (tokenized.title && this.exactMatch(tokenized.title, token)) { + score += this.SCORING_WEIGHTS.TITLE; + } + } + return { score }; + } + async getVaultEntries(): Promise { const files = this.vault.getMarkdownFiles(); const entries: VaultEntry[] = []; for (const file of files) { try { - const content = await this.vault.cachedRead(file); + const content = + typeof this.vault.cachedRead === 'function' + ? await this.vault.cachedRead(file) + : await this.vault.read(file); const parsed = this.parseMarkdown(content); entries.push({ file: file, - title: parsed.title, + title: parsed.frontmatter.title || file.basename, frontmatter: parsed.frontmatter, headings: parsed.headings, - content: parsed.content, + content: parsed.content.slice(0, CONTENT_PREVIEW_LENGTH), basename: file.basename, + score: 0, }); } catch (error) { const errorMessage = error instanceof Error ? error.message : String(error); @@ -76,7 +161,6 @@ export class VaultIndexer { try { cachedResults = await this.cache.get(cacheKey); } catch { - // Ignore cache retrieval errors and continue with normal processing cachedResults = null; } if (cachedResults) { @@ -85,53 +169,36 @@ export class VaultIndexer { const parsedResults: VaultEntry[] = JSON.parse(cachedResults); return parsedResults.slice(0, limit); } catch { - // Ignore cache parse errors and continue with normal processing + // ignore parse errors } } } - const entries = await this.getVaultEntries(); - const tokenized = this.tokenizeQuery(query); - const scoredEntries = entries.map((entry) => { - let score = 0; - for (const queryToken of tokenized.tokens) { - const stemmed = this.stemToken(queryToken); - let matched = false; - if (entry.frontmatter?.title && this.exactMatch(entry.frontmatter.title, queryToken)) { - score += this.SCORING_WEIGHTS.FRONTMATTER_TITLE; - matched = true; - } else if ( - entry.basename && - this.exactMatch(entry.basename.replace(/\.md$/, ''), queryToken) - ) { - score += this.SCORING_WEIGHTS.FRONTMATTER_TITLE; - matched = true; - } - if (entry.frontmatter?.tags && this.exactMatch(entry.frontmatter.tags, queryToken)) { - score += this.SCORING_WEIGHTS.FRONTMATTER_TAGS; - matched = true; - } - if (entry.headings.some((heading) => heading.toLowerCase().includes(stemmed))) { - score += this.SCORING_WEIGHTS.HEADINGS; - matched = true; - } - if (entry.content.toLowerCase().includes(stemmed)) { - score += this.SCORING_WEIGHTS.CONTENT; - matched = true; - } - if (matched) { - // Add bonus for exact matches - if (entry.title && this.exactMatch(entry.title, queryToken)) { - score += this.SCORING_WEIGHTS.TITLE; - } - } - } - return { entry, score }; - }); + const queryTokens = this.tokenize(query); + if (queryTokens.length === 0) { + return []; + } - // Sort by score descending and return top results - scoredEntries.sort((a, b) => b.score - a.score); - const results = scoredEntries.slice(0, limit).map((item) => item.entry); + const entries = await this.getVaultEntries(); + const scored = entries + .map((entry) => { + const { score } = this.calculateWeightedScore( + { + title: entry.title, + headings: entry.headings, + frontmatter: entry.frontmatter, + firstParagraph: '', + content: entry.content, + basename: entry.basename, + }, + queryTokens + ); + return { ...entry, score }; + }) + .filter((e) => e.score > 0); + + scored.sort((a, b) => b.score - a.score); + const results = scored.slice(0, limit); if (this.cache) { try { @@ -148,49 +215,34 @@ export class VaultIndexer { return results; } - private tokenizeQuery(query: string) { - const tokens = query - .toLowerCase() - .replace(/[^\w\s]/g, '') - .split(/\s+/) - .filter((token) => token.length > 0); - return { - tokens, - stemmedTokens: tokens.map((token) => this.stemToken(token)), - }; - } - - private stemToken(token: string) { - // Simple stemming for now - in a real implementation, this would be more sophisticated + private stemToken(token: string): string { + if (token.endsWith('ing') && token.length > 4) return token.slice(0, -3); + if (token.endsWith('ed') && token.length > 3) return token.slice(0, -2); + if (token.endsWith('s') && token.length > 2) return token.slice(0, -1); return token; } - private exactMatch(text: string | undefined, queryToken: string) { + private exactMatch(text: string | undefined, queryToken: string): boolean { if (!text) return false; - return text.toLowerCase().includes(queryToken.toLowerCase()); + const textLower = text.toLowerCase(); + const queryLower = queryToken.toLowerCase(); + const queryStem = this.stemToken(queryLower); + return textLower.includes(queryLower) || textLower.includes(queryStem); } - private parseMarkdown(content: string): ParsedMarkdown { + private parseMarkdown(content: string) { const frontmatterRegex = /^---\n([\s\S]*?)\n---\n/; const frontmatterMatch = content.match(frontmatterRegex); const frontmatter: ParsedFrontmatter = {}; if (frontmatterMatch) { try { - const frontmatterContent = frontmatterMatch[1]; - const lines = frontmatterContent.trim().split('\n'); + const lines = frontmatterMatch[1].trim().split('\n'); for (const line of lines) { const [key, ...valueParts] = line.split(':'); if (!key) continue; const value = valueParts.join(':').trim(); - if (key.trim() === 'title') { - if (value) { - frontmatter.title = value; - } - } else if (key.trim() === 'tags') { - if (value) { - frontmatter.tags = value; - } - } + if (key.trim() === 'title' && value) frontmatter.title = value; + else if (key.trim() === 'tags' && value) frontmatter.tags = value; } } catch { Logger.warn('Failed to parse frontmatter', 'vault-indexer'); @@ -207,20 +259,14 @@ export class VaultIndexer { headings.push(headingMatch[1]); } - const contentWithoutFrontmatter = frontmatterMatch + const bodyWithoutFrontmatter = frontmatterMatch ? content.substring(frontmatterMatch[0].length) : content; - const contentWithoutFrontmatterAndHeadings = contentWithoutFrontmatter - .replace(/#{1,6} .+/g, '') // Remove headings - .replace(/^---[\s\S]*?---\n/, '') // Remove frontmatter - .replace(/^\s*[\r\n]/gm, '') // Remove empty lines + const bodyText = bodyWithoutFrontmatter + .replace(/#{1,6} .+/g, '') + .replace(/^\s*[\r\n]/gm, '') .trim(); - return { - frontmatter, - title, - headings, - content: contentWithoutFrontmatterAndHeadings, - }; + return { frontmatter, title, headings, content: bodyText }; } } diff --git a/tests/ollama-client-cache.test.ts b/tests/ollama-client-cache.test.ts index 7f5ee22..542fb28 100644 --- a/tests/ollama-client-cache.test.ts +++ b/tests/ollama-client-cache.test.ts @@ -16,7 +16,7 @@ jest.mock('../src/semantic-cache', () => ({ setCache: mockSetCache, clearCache: mockClearCache, })), -)); +})); import { OllamaClient } from '../src/ollama-client'; diff --git a/tests/semantic-cache.test.ts b/tests/semantic-cache.test.ts index 9558910..3d12f69 100644 --- a/tests/semantic-cache.test.ts +++ b/tests/semantic-cache.test.ts @@ -39,16 +39,21 @@ describe('SemanticCacheService', () => { let mockChromaClient: any; let mockCollection: any; - beforeEach(() => { + beforeEach(async () => { // Reset all mocks jest.clearAllMocks(); // Create a fresh instance for each test cacheService = new SemanticCacheService(mockOllamaUrl, mockCacheConfig); + await cacheService.initialize(); + global.fetch = jest.fn().mockResolvedValue({ + ok: true, + json: () => Promise.resolve({ embedding: [0.1, 0.2, 0.3] }), + }); // Access the internal mocks - mockChromaClient = (ChromaClient as jest.Mock).mock.instances[0]; - mockCollection = mockChromaClient.getOrCreateCollection.mock.results[0].value; + mockChromaClient = (ChromaClient as jest.Mock).mock.results[0].value; + mockCollection = await mockChromaClient.getOrCreateCollection.mock.results[0].value; }); describe('constructor', () => { @@ -72,6 +77,7 @@ describe('SemanticCacheService', () => { }); it('should not initialize when cache is disabled', async () => { + jest.clearAllMocks(); const disabledConfig: CacheConfig = { ...mockCacheConfig, enabled: false }; const disabledCacheService = new SemanticCacheService(mockOllamaUrl, disabledConfig); @@ -83,6 +89,7 @@ describe('SemanticCacheService', () => { describe('getCache', () => { it('should return null when cache is disabled', async () => { + jest.clearAllMocks(); const disabledConfig: CacheConfig = { ...mockCacheConfig, enabled: false }; const disabledCacheService = new SemanticCacheService(mockOllamaUrl, disabledConfig); @@ -122,6 +129,7 @@ describe('SemanticCacheService', () => { describe('setCache', () => { it('should not set cache when disabled', async () => { + jest.clearAllMocks(); const disabledConfig: CacheConfig = { ...mockCacheConfig, enabled: false }; const disabledCacheService = new SemanticCacheService(mockOllamaUrl, disabledConfig); @@ -131,25 +139,15 @@ describe('SemanticCacheService', () => { }); it('should add content to cache', async () => { - const mockEmbedding = [0.1, 0.2, 0.3]; - - // Mock the fetch function for embedding generation - global.fetch = jest.fn().mockResolvedValue({ - ok: true, - json: () => Promise.resolve({ embedding: mockEmbedding }), - }); - await cacheService.setCache('test query', 'test response'); expect(mockCollection.add).toHaveBeenCalled(); - - // Clean up - global.fetch = undefined as any; }); }); describe('clearCache', () => { it('should not clear when cache is disabled', async () => { + jest.clearAllMocks(); const disabledConfig: CacheConfig = { ...mockCacheConfig, enabled: false }; const disabledCacheService = new SemanticCacheService(mockOllamaUrl, disabledConfig);