From 771db09d2435e955399fc49c1ef5121f4814038e Mon Sep 17 00:00:00 2001 From: Florian Egger Date: Fri, 8 May 2026 11:52:31 +0200 Subject: [PATCH] Refactor chat view and add conversation state management Introduce ConversationStateManager to handle short, medium, and long-term context for improved conversation flow. Update ChatView to use this manager and refactor input handling to accept values directly for better testability. Update OllamaClient with non-streaming chat support and improved error handling for malformed chunks. Enhance vault indexer with caching, better scoring, and stop word filtering. Refactor main plugin entry point and semantic cache initialization for robustness. --- src/cache.js | 2 + src/chat-view.js | 859 +++++++++++++------------ src/chat-view.ts | 173 ++--- src/constants.js | 29 +- src/conversation-state.js | 145 +++++ src/conversation-state.ts | 171 +++++ src/error-handler.js | 2 +- src/error-handler.ts | 2 +- src/indexing-pipeline/extraction.js | 4 +- src/indexing-pipeline/extraction.ts | 7 +- src/indexing-pipeline/normalization.ts | 7 - src/indexing-pipeline/pipeline.js | 15 +- src/indexing-pipeline/pipeline.ts | 24 +- src/indexing-pipeline/vectorization.js | 12 +- src/indexing-pipeline/vectorization.ts | 18 +- src/main.js | 374 ++++++----- src/main.ts | 7 +- src/ollama-client.js | 451 ++++++++----- src/ollama-client.ts | 205 +++++- src/semantic-cache.js | 179 +++--- src/tool-executor.js | 53 ++ src/tool-executor.ts | 64 +- src/types.js | 4 +- src/types.ts | 53 +- src/utils.js | 3 +- src/utils.ts | 5 +- src/vault-indexer.js | 355 ++++++---- src/vault-indexer.ts | 220 ++++--- tests/ollama-client-cache.test.ts | 2 +- tests/semantic-cache.test.ts | 26 +- 30 files changed, 2184 insertions(+), 1287 deletions(-) create mode 100644 src/cache.js create mode 100644 src/conversation-state.js create mode 100644 src/conversation-state.ts 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);