diff --git a/src/chat-view.js b/src/chat-view.js index e3cbd8a..7d352a9 100644 --- a/src/chat-view.js +++ b/src/chat-view.js @@ -1,423 +1,452 @@ -"use strict"; -Object.defineProperty(exports, "__esModule", { value: true }); +'use strict'; +Object.defineProperty(exports, '__esModule', { value: true }); exports.ChatView = void 0; -const obsidian_1 = require("obsidian"); +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"); +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'); class ChatView extends obsidian_1.ItemView { - // Getters for testing - getSendButtonClickHandler() { - return this.sendButtonClickHandler; + // 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.' + ); } - getInputKeyDownHandler() { - return this.inputKeyDownHandler; + 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; } - getNewChatButtonClickHandler() { - return this.newChatButtonClickHandler; + } + 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(); + } } - 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); + // 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; + } } - 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.'); - }); + // 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); + } } - async clearCache() { - await this.ollamaClient.clearCache(); + // 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); } - getViewType() { - return 'ollama-chat-view'; + // 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); } - getDisplayText() { - return 'Ollama Chat'; + // 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); } - 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.'); - } + // 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 }; this.render(); - this.removeEventListeners(); // Clean up any existing listeners before reattaching - this.setupEventListeners(); + } } - onSettingsChange(newSettings) { - this.updateSettings(newSettings); - } - 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' }); - if (!this.inputEl) { - this.inputEl = inputContainer.createEl('textarea', { cls: 'ollama-input' }); - } - if (!this.sendButton) { - this.sendButton = inputContainer.createEl('button', { - cls: 'ollama-send-button', - }); - this.sendButton.textContent = 'Send'; - } - if (!this.newChatButton) { - const newChatContainer = this.contentEl.querySelector('.ollama-new-chat') || - this.contentEl.createEl('div', { cls: 'ollama-new-chat' }); - this.newChatButton = newChatContainer.createEl('button', { - cls: 'ollama-new-chat-button', - }); - this.newChatButton.textContent = '🔄 New Chat'; - this.newChatButton.title = 'Start a new conversation'; - } - // Create immutable snapshot for rendering - const messagesSnapshot = [...this.messages]; - // Only render messages that are not currently streaming - const nonStreamingMessages = messagesSnapshot.filter((msg) => !msg.isStreaming); - // Differential update: only update messages that have changed - const existingMessages = container.querySelectorAll('.ollama-message'); - for (const msg of nonStreamingMessages) { - const existingEl = container.querySelector(`.ollama-message[data-msg-id="${msg.id}"]`); - if (existingEl) { - existingEl.textContent = msg.content; - } - else { - const messageEl = container.createEl('div', { - cls: `ollama-message ${msg.role}`, - }); - messageEl.setAttribute('data-msg-id', msg.id); - messageEl.textContent = msg.content; - } - } - // 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(); - } - } - // 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); - } - } - } - setupEventListeners() { - if (!this.sendButton || !this.inputEl || this.listenersAttached) - return; - // Create handlers if they don't exist - if (!this.sendButtonClickHandler) { - this.sendButtonClickHandler = async () => { - if (!this.inputEl) - return; - await this.handleUserInput(this.inputEl.value); - this.inputEl.value = ''; - }; - } - if (!this.inputKeyDownHandler) { - this.inputKeyDownHandler = async (e) => { - if (!this.inputEl || e.key !== 'Enter' || e.shiftKey) - return; - e.preventDefault(); - await this.handleUserInput(this.inputEl.value); - this.inputEl.value = ''; - }; - } - // Create wrapper functions for event listeners - this.sendButtonClickWrapper = () => { - void this.sendButtonClickHandler?.(); - }; - this.inputKeyDownWrapper = (e) => { - void this.inputKeyDownHandler?.(e); - }; - this.newChatButtonClickWrapper = () => { - void this.newChatButtonClickHandler?.(); - }; - // Add event listeners using wrappers - this.sendButton.addEventListener('click', this.sendButtonClickWrapper); - this.inputEl.addEventListener('keydown', this.inputKeyDownWrapper); - if (this.newChatButton) { - if (!this.newChatButtonClickHandler) { - this.newChatButtonClickHandler = () => this.clearConversation(); - } - this.newChatButton.addEventListener('click', this.newChatButtonClickWrapper); - } - this.listenersAttached = true; - } - removeEventListeners() { - if (this.sendButton && this.sendButtonClickWrapper) { - this.sendButton.removeEventListener('click', this.sendButtonClickWrapper); - } - if (this.inputEl && this.inputKeyDownWrapper) { - this.inputEl.removeEventListener('keydown', this.inputKeyDownWrapper); - } - if (this.newChatButton && this.newChatButtonClickWrapper) { - this.newChatButton.removeEventListener('click', this.newChatButtonClickWrapper); - } - this.sendButtonClickWrapper = null; - this.inputKeyDownWrapper = null; - this.newChatButtonClickWrapper = null; - this.listenersAttached = false; - } - clearConversation() { - // Create new array to ensure immutability - this.messages = []; - this.lastMessageEl = null; - this.render(); - new obsidian_1.Notice('Conversation cleared'); - } - updateMessageById(id, partial) { - const index = this.messages.findIndex((m) => m.id === id); - if (index < 0) - return false; - this.messages = [ - ...this.messages.slice(0, index), - { ...this.messages[index], ...partial }, - ...this.messages.slice(index + 1), - ]; - return true; - } - updateLastMessage(content) { - const streamingMessage = this.messages.find((msg) => msg.isStreaming); - if (streamingMessage && !this.lastMessageEl) { - this.lastMessageEl = (this.chatContainer ?? this.contentEl).createEl('div', { - cls: `ollama-message assistant`, - }); - this.lastMessageEl.setAttribute('data-msg-id', streamingMessage.id); - } - if (this.lastMessageEl) { - this.lastMessageEl.textContent = content; - } - } - getTools() { - return [ - { - type: 'function', - function: { - name: 'create_file', - description: 'Create a new file in the vault', - parameters: { - type: 'object', - properties: { - path: { - type: 'string', - description: "Relative path within the vault, e.g. 'Notes/todo.md'", - }, - content: { type: 'string', description: 'Content of the file to create' }, - }, - required: ['path', 'content'], - }, - }, + } + 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'], + }, + }, + }, + { + 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'], + }, + }, + }, + ]; + } + buildMessages(userMessageWithContext, 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, + }); } - buildMessages(userMessage, context) { - const systemContent = context - ? `You are a helpful assistant.\n\nRelevant vault context:\n${context}` - : 'You are a helpful assistant.'; - const systemMessage = { - role: 'system', - content: systemContent, - }; - const userMessageWithContext = { - role: 'user', - content: userMessage, - }; - return [ - systemMessage, - ...this.messages.map((m) => ({ - role: m.role, - content: m.content, - tool_calls: m.tool_calls, - })), - userMessageWithContext, - ]; - } - async processToolCalls(toolCalls, messages, tools, fullResponse, assistantMessageId) { - // Validate tool calls before processing - const MAX_TOOL_CALLS = 10; - if (toolCalls.length > MAX_TOOL_CALLS) { - throw new Error(`Too many tool calls (max ${MAX_TOOL_CALLS})`); - } - // Collect all tool results using allSettled to support partial results - const settledResults = await Promise.allSettled(toolCalls.map((call) => this.toolExecutor.handleToolCall(call))); - const toolResults = []; - for (const result of settledResults) { - if (result.status === 'fulfilled') { - toolResults.push(result.value); - } - else { - // Use centralized error handler for tool errors - error_handler_1.ErrorHandler.handleError(result.reason, 'ChatView.handleUserInput'); - } - } - // Only create follow-up when we have tool results - if (toolResults.length > 0) { - // Create follow-up messages including the assistant's tool calls and results - const followUpMessages = [ - ...messages, - { role: 'assistant', content: fullResponse, tool_calls: toolCalls }, - ...toolResults.map((result) => ({ - role: 'tool', - content: JSON.stringify(result), - })), - ]; - const followUp = await this.ollamaClient.chat(followUpMessages, tools); - fullResponse += followUp.content; - this.updateLastMessage(fullResponse); - // Update the assistant message with the final response immutably - this.updateMessageById(assistantMessageId, { - content: fullResponse, - isStreaming: false, - }); - } - else { - // Even if no tool results were successful, mark streaming as complete - // to prevent the assistant message from disappearing - this.updateMessageById(assistantMessageId, { - content: fullResponse, - isStreaming: false, - }); - } - } - async handleUserInput(content) { - if (!this.sendButton || !this.inputEl) - return; - this.sendButton.disabled = true; + return messages; + } + async processToolCalls(toolCalls, messages, tools, fullResponse, assistantMessageId) { + const MAX_TOOL_CALLS = 5; + const settledResults = await Promise.allSettled( + toolCalls.map(async (toolCall) => { try { - // Guard against empty messages - const userMessage = content.trim(); - if (!userMessage) - return; - // Search vault using user message as query - const entries = await this.vaultIndexer.searchVault(userMessage, this.settings.vaultSearchLimit); - let context = entries.map((entry) => `### ${entry.title}\n${entry.content}`).join('\n\n'); - // Cap context size to prevent prompt bloat with large vaults - const MAX_CONTEXT_LENGTH = 4000; - if (context.length > MAX_CONTEXT_LENGTH) { - context = context.substring(0, MAX_CONTEXT_LENGTH) + '\n\n... (truncated)'; - } - const messages = this.buildMessages(userMessage, context); - const tools = this.getTools(); - const messageId = `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; - const userMessageId = messageId; - const assistantMessageId = `${messageId}-assistant`; - // Store user message in conversation history - const userChatMessage = { - id: userMessageId, - role: 'user', - content: userMessage, - timestamp: Date.now(), - }; - const assistantMessage = { - id: assistantMessageId, - role: 'assistant', - content: '', - timestamp: Date.now(), - isStreaming: true, - }; - // Update messages immutably - this.messages = [...this.messages, userChatMessage, assistantMessage]; - try { - this.render(); - const stream = this.ollamaClient.streamChat(messages, tools); - let fullResponse = ''; - let toolCalls = []; - let chunkCount = 0; - for await (const chunk of stream) { - chunkCount++; - if (chunkCount > MAX_STREAM_CHUNKS) { - throw new Error('Response too long, stopped streaming'); - } - if (chunk.content) { - fullResponse += chunk.content; - } - if (chunk.tool_calls) { - toolCalls = toolCalls.concat(chunk.tool_calls); - } - this.updateLastMessage(fullResponse); - } - // Update the assistant message with the full response immutably - this.updateMessageById(assistantMessageId, { - content: fullResponse, - tool_calls: toolCalls, - }); - // Process tool calls with proper follow-up context - if (toolCalls.length > 0) { - await this.processToolCalls(toolCalls, messages, 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(); - } - } - catch (error) { - // Use centralized error handler - error_handler_1.ErrorHandler.handleError(error, 'ChatView.handleUserInput'); - // Update any streaming messages to non-streaming state to prevent stale messages - this.messages = this.messages.map((msg) => msg.isStreaming ? { ...msg, isStreaming: false } : msg); - this.cleanupStreamingResources(); - this.render(); - } - finally { - if (this.sendButton) { - this.sendButton.disabled = false; - } + 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 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, + }); + } + this.updateMessageById(assistantMessageId, { + content: fullResponse, + 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, + isStreaming: true, + }); + } + if (chunk.message.tool_calls) { + toolCalls = [...toolCalls, ...chunk.message.tool_calls]; + } + chunkCount++; + if (chunkCount > MAX_STREAM_CHUNKS) { + break; + } + } + // 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; diff --git a/src/chat-view.ts b/src/chat-view.ts index 2e4403d..07add57 100755 --- a/src/chat-view.ts +++ b/src/chat-view.ts @@ -1,58 +1,41 @@ -import { ItemView, WorkspaceLeaf, Notice } from 'obsidian'; -/// -// Use global types from JSDOM setup -type KeyboardEvent = globalThis.KeyboardEvent; -type HTMLTextAreaElement = globalThis.HTMLTextAreaElement; -type HTMLButtonElement = globalThis.HTMLButtonElement; - -const MAX_STREAM_CHUNKS = 1000; -import { - PluginSettings, - OllamaMessage, - ChatMessage, - OllamaTool, - ToolCall, - ToolResult, -} from './types'; +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 { ErrorHandler } from './error-handler'; -export class ChatView extends ItemView { - private settings: PluginSettings; - private messages: ChatMessage[] = []; - private ollamaClient: OllamaClient; - private vaultIndexer: VaultIndexer; - private toolExecutor: ToolExecutor; - private lastMessageEl: HTMLElement | null = null; - private newChatButton: HTMLButtonElement | null = null; - private sendButton: HTMLButtonElement | null = null; - private inputEl: HTMLTextAreaElement | null = null; - private chatContainer: HTMLElement | null = null; - private sendButtonClickHandler: (() => Promise) | null = null; - private inputKeyDownHandler: ((e: KeyboardEvent) => Promise) | null = null; - private newChatButtonClickHandler: (() => void) | null = null; - private sendButtonClickWrapper: (() => void) | null = null; - private inputKeyDownWrapper: ((e: KeyboardEvent) => void) | null = null; - private newChatButtonClickWrapper: (() => void) | null = null; - private listenersAttached = false; +export const VIEW_TYPE_OLLAMA_CHAT = 'ollama-chat-view'; +export class ChatView extends ItemView { // Getters for testing - public getSendButtonClickHandler(): (() => Promise) | null { + getSendButtonClickHandler() { return this.sendButtonClickHandler; } - public getInputKeyDownHandler(): ((e: KeyboardEvent) => Promise) | null { + getInputKeyDownHandler() { return this.inputKeyDownHandler; } - public getNewChatButtonClickHandler(): (() => void) | null { + getNewChatButtonClickHandler() { return this.newChatButtonClickHandler; } constructor(leaf: WorkspaceLeaf, settings: PluginSettings) { 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 OllamaClient( settings.ollamaUrl, @@ -64,7 +47,7 @@ export class ChatView extends ItemView { this.toolExecutor = new ToolExecutor(this.app.vault, this.app); } - public updateSettings(newSettings: PluginSettings): void { + updateSettings(newSettings: PluginSettings) { this.settings = newSettings; this.ollamaClient = new OllamaClient( newSettings.ollamaUrl, @@ -79,7 +62,7 @@ export class ChatView extends ItemView { }); } - public async clearCache(): Promise { + async clearCache(): Promise { await this.ollamaClient.clearCache(); } @@ -104,11 +87,11 @@ export class ChatView extends ItemView { this.setupEventListeners(); } - public onSettingsChange(newSettings: PluginSettings): void { + onSettingsChange(newSettings: PluginSettings): void { this.updateSettings(newSettings); } - onClose(): Promise { + async onClose(): Promise { this.ollamaClient.cancelStream(); this.removeEventListeners(); this.cleanupStreamingResources(); @@ -119,7 +102,7 @@ export class ChatView extends ItemView { return Promise.resolve(); } - private cleanupStreamingResources(): void { + cleanupStreamingResources(): void { // 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) { @@ -128,57 +111,21 @@ export class ChatView extends ItemView { } } - render() { + render(): void { 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' }); - if (!this.inputEl) { - this.inputEl = inputContainer.createEl('textarea', { cls: 'ollama-input' }); - } - if (!this.sendButton) { - this.sendButton = inputContainer.createEl('button', { - cls: 'ollama-send-button', - }); - this.sendButton.textContent = 'Send'; - } - - if (!this.newChatButton) { - const newChatContainer = - this.contentEl.querySelector('.ollama-new-chat') || - this.contentEl.createEl('div', { cls: 'ollama-new-chat' }); - this.newChatButton = newChatContainer.createEl('button', { - cls: 'ollama-new-chat-button', - }); - this.newChatButton.textContent = '🔄 New Chat'; - this.newChatButton.title = 'Start a new conversation'; - } - - // Create immutable snapshot for rendering const messagesSnapshot = [...this.messages]; - - // Only render messages that are not currently streaming const nonStreamingMessages = messagesSnapshot.filter((msg) => !msg.isStreaming); - - // Differential update: only update messages that have changed const existingMessages = container.querySelectorAll('.ollama-message'); - for (const msg of nonStreamingMessages) { - const existingEl = container.querySelector(`.ollama-message[data-msg-id="${msg.id}"]`); - if (existingEl) { - existingEl.textContent = msg.content; - } else { - const messageEl = container.createEl('div', { - cls: `ollama-message ${msg.role}`, - }) as HTMLElement; - messageEl.setAttribute('data-msg-id', msg.id); - messageEl.textContent = msg.content; - } - } - // Remove messages that are no longer in the array for (const el of Array.from(existingMessages)) { const id = el.getAttribute('data-msg-id'); @@ -187,6 +134,20 @@ export class ChatView extends ItemView { } } + // 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; + } + } + // Re-attach streaming message if it exists const streamingMessage = messagesSnapshot.find((msg) => msg.isStreaming); if (streamingMessage && this.lastMessageEl) { @@ -197,199 +158,263 @@ export class ChatView extends ItemView { container.appendChild(this.lastMessageEl); } } + + // 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); + } + + // 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); + } + + // 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); + } + + // Append containers to contentEl + this.contentEl.appendChild(newChatContainer); + this.contentEl.appendChild(inputContainer); + this.contentEl.appendChild(container); + + // Focus input on open + this.inputEl.focus(); } - private setupEventListeners(): void { - if (!this.sendButton || !this.inputEl || this.listenersAttached) return; - - // Create handlers if they don't exist - if (!this.sendButtonClickHandler) { - this.sendButtonClickHandler = async () => { - if (!this.inputEl) return; - await this.handleUserInput(this.inputEl.value); - this.inputEl.value = ''; - }; + setupEventListeners(): void { + if (this.listenersAttached) { + return; } - if (!this.inputKeyDownHandler) { - this.inputKeyDownHandler = async (e: KeyboardEvent) => { - if (!this.inputEl || e.key !== 'Enter' || e.shiftKey) return; - e.preventDefault(); - await this.handleUserInput(this.inputEl.value); - this.inputEl.value = ''; - }; - } - - // Create wrapper functions for event listeners - this.sendButtonClickWrapper = () => { - void this.sendButtonClickHandler?.(); - }; - this.inputKeyDownWrapper = (e: KeyboardEvent) => { - void this.inputKeyDownHandler?.(e); - }; - this.newChatButtonClickWrapper = () => { - void this.newChatButtonClickHandler?.(); + this.sendButtonClickHandler = () => { + void this.handleUserInput(); }; - // Add event listeners using wrappers - this.sendButton.addEventListener('click', this.sendButtonClickWrapper); - this.inputEl.addEventListener('keydown', this.inputKeyDownWrapper); - if (this.newChatButton) { - if (!this.newChatButtonClickHandler) { - this.newChatButtonClickHandler = () => this.clearConversation(); + this.inputKeyDownHandler = (event: KeyboardEvent) => { + if (event.key === 'Enter' && !event.shiftKey) { + event.preventDefault(); + void this.handleUserInput(); } - this.newChatButton.addEventListener('click', this.newChatButtonClickWrapper); + }; + + 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; } - private removeEventListeners(): void { - if (this.sendButton && this.sendButtonClickWrapper) { - this.sendButton.removeEventListener('click', this.sendButtonClickWrapper); + removeEventListeners(): void { + if (!this.listenersAttached) { + return; } - if (this.inputEl && this.inputKeyDownWrapper) { - this.inputEl.removeEventListener('keydown', this.inputKeyDownWrapper); + + if (this.sendButton) { + this.sendButton.removeEventListener('click', this.sendButtonClickHandler); } - if (this.newChatButton && this.newChatButtonClickWrapper) { - this.newChatButton.removeEventListener('click', this.newChatButtonClickWrapper); + + if (this.inputEl) { + this.inputEl.removeEventListener('keydown', this.inputKeyDownHandler); } - this.sendButtonClickWrapper = null; - this.inputKeyDownWrapper = null; - this.newChatButtonClickWrapper = null; + + if (this.newChatButton) { + this.newChatButton.removeEventListener('click', this.newChatButtonClickHandler); + } + this.listenersAttached = false; } - private clearConversation(): void { - // Create new array to ensure immutability + clearConversation(): void { this.messages = []; - this.lastMessageEl = null; this.render(); - new Notice('Conversation cleared'); } - private updateMessageById(id: string, partial: Partial): boolean { + updateMessageById(id: string, updates: Partial): void { const index = this.messages.findIndex((m) => m.id === id); - if (index < 0) return false; - this.messages = [ - ...this.messages.slice(0, index), - { ...this.messages[index], ...partial }, - ...this.messages.slice(index + 1), - ]; - return true; + if (index !== -1) { + this.messages[index] = { ...this.messages[index], ...updates }; + this.render(); + } } - private updateLastMessage(content: string) { + updateLastMessage(updates: Partial): void { const streamingMessage = this.messages.find((msg) => msg.isStreaming); - if (streamingMessage && !this.lastMessageEl) { - this.lastMessageEl = (this.chatContainer ?? this.contentEl).createEl('div', { - cls: `ollama-message assistant`, - }); - this.lastMessageEl.setAttribute('data-msg-id', streamingMessage.id); - } - if (this.lastMessageEl) { - this.lastMessageEl.textContent = content; + if (streamingMessage) { + const index = this.messages.findIndex((msg) => msg.id === streamingMessage.id); + if (index !== -1) { + this.messages[index] = { ...this.messages[index], ...updates }; + this.render(); + } } } - private getTools(): OllamaTool[] { + getTools(): OllamaTool[] { return [ { type: 'function', function: { - name: 'create_file', - description: 'Create a new file in the vault', + name: 'read_vault_file', + description: 'Reads the content of a file from the vault', parameters: { - type: 'object' as const, + type: 'object', properties: { path: { - type: 'string' as const, - description: "Relative path within the vault, e.g. 'Notes/todo.md'", + type: 'string', + description: 'The path to the file to read', + }, + content: { + type: 'string', + description: 'The content of the file to read', }, - content: { type: 'string' as const, description: 'Content of the file to create' }, }, - required: ['path', 'content'], + 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', + }, + }, + required: ['query'], }, }, }, ]; } - private buildMessages(userMessage: string, context: string): OllamaMessage[] { - const systemContent = context - ? `You are a helpful assistant.\n\nRelevant vault context:\n${context}` - : 'You are a helpful assistant.'; - const systemMessage: OllamaMessage = { + buildMessages(userMessageWithContext: 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 = { role: 'system', content: systemContent, }; - const userMessageWithContext: OllamaMessage = { + + const userMessageWithContext = { role: 'user', - content: userMessage, + content: userMessageWithContext, }; - return [ - systemMessage, - ...this.messages.map((m) => ({ - role: m.role, - content: m.content, - tool_calls: m.tool_calls, - })), - userMessageWithContext, - ]; + const messages: OllamaMessage[] = [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; } - private async processToolCalls( - toolCalls: ToolCall[], + async processToolCalls( + toolCalls: OllamaToolCall[], messages: OllamaMessage[], tools: OllamaTool[], fullResponse: string, assistantMessageId: string ): Promise { - // Validate tool calls before processing - const MAX_TOOL_CALLS = 10; - if (toolCalls.length > MAX_TOOL_CALLS) { - throw new Error(`Too many tool calls (max ${MAX_TOOL_CALLS})`); - } - - // Collect all tool results using allSettled to support partial results + const MAX_TOOL_CALLS = 5; const settledResults = await Promise.allSettled( - toolCalls.map((call) => this.toolExecutor.handleToolCall(call)) + 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 toolResults: ToolResult[] = []; - for (const result of settledResults) { - if (result.status === 'fulfilled') { - toolResults.push(result.value); - } else { - // Use centralized error handler for tool errors - ErrorHandler.handleError(result.reason, 'ChatView.handleUserInput'); + 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, + }); } - } - - // Only create follow-up when we have tool results - if (toolResults.length > 0) { - // Create follow-up messages including the assistant's tool calls and results - const followUpMessages: OllamaMessage[] = [ - ...messages, - { role: 'assistant' as const, content: fullResponse, tool_calls: toolCalls }, - ...toolResults.map((result) => ({ - role: 'tool' as const, - content: JSON.stringify(result), - })), - ]; - - const followUp = await this.ollamaClient.chat(followUpMessages, tools); - fullResponse += followUp.content; - this.updateLastMessage(fullResponse); - - // Update the assistant message with the final response immutably - this.updateMessageById(assistantMessageId, { - content: fullResponse, - isStreaming: false, - }); - } else { - // Even if no tool results were successful, mark streaming as complete - // to prevent the assistant message from disappearing this.updateMessageById(assistantMessageId, { content: fullResponse, isStreaming: false, @@ -397,119 +422,131 @@ export class ChatView extends ItemView { } } - private async handleUserInput(content: string) { - if (!this.sendButton || !this.inputEl) return; - this.sendButton.disabled = true; + async handleUserInput(): Promise { + 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 { - // Guard against empty messages - const userMessage = content.trim(); - if (!userMessage) return; + const stream = await this.ollamaClient.streamChat(this.buildMessages(userMessage, tools), { + temperature: 0.5, + }); - // Search vault using user message as query - const entries = await this.vaultIndexer.searchVault( - userMessage, - this.settings.vaultSearchLimit - ); - let context = entries.map((entry) => `### ${entry.title}\n${entry.content}`).join('\n\n'); + let fullResponse = ''; + let toolCalls: OllamaToolCall[] = []; + let chunkCount = 0; - // Cap context size to prevent prompt bloat with large vaults - const MAX_CONTEXT_LENGTH = 4000; - if (context.length > MAX_CONTEXT_LENGTH) { - context = context.substring(0, MAX_CONTEXT_LENGTH) + '\n\n... (truncated)'; - } - - const messages = this.buildMessages(userMessage, context); - const tools = this.getTools(); - - const messageId = `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`; - - const userMessageId = messageId; - const assistantMessageId = `${messageId}-assistant`; - - // Store user message in conversation history - const userChatMessage: ChatMessage = { - id: userMessageId, - role: 'user' as const, - content: userMessage, - timestamp: Date.now(), - }; - - const assistantMessage: ChatMessage = { - id: assistantMessageId, - role: 'assistant' as const, - content: '', - timestamp: Date.now(), - isStreaming: true, - }; - - // Update messages immutably - this.messages = [...this.messages, userChatMessage, assistantMessage]; - - try { - this.render(); - - const stream = this.ollamaClient.streamChat(messages, tools); - let fullResponse = ''; - let toolCalls: ToolCall[] = []; - let chunkCount = 0; - for await (const chunk of stream) { - chunkCount++; - if (chunkCount > MAX_STREAM_CHUNKS) { - throw new Error('Response too long, stopped streaming'); - } - - if (chunk.content) { - fullResponse += chunk.content; - } - - if (chunk.tool_calls) { - toolCalls = toolCalls.concat(chunk.tool_calls); - } - - this.updateLastMessage(fullResponse); - } - - // Update the assistant message with the full response immutably - this.updateMessageById(assistantMessageId, { - content: fullResponse, - tool_calls: toolCalls, - }); - - // Process tool calls with proper follow-up context - if (toolCalls.length > 0) { - await this.processToolCalls(toolCalls, messages, tools, fullResponse, assistantMessageId); - } - - // Update assistant message immutably — only if no tool calls were processed - if (toolCalls.length === 0) { - this.updateMessageById(assistantMessageId, { - isStreaming: false, + for await (const chunk of stream) { + if (chunk.message.content) { + fullResponse += chunk.message.content; + this.updateLastMessage({ + content: fullResponse, + isStreaming: true, }); } - // Limit conversation history to prevent memory issues - if (this.messages.length > this.settings.maxMessageHistory) { - this.messages = this.messages.slice(-this.settings.maxMessageHistory); + if (chunk.message.tool_calls) { + toolCalls = [...toolCalls, ...chunk.message.tool_calls]; + } + + chunkCount++; + if (chunkCount > MAX_STREAM_CHUNKS) { + break; } - this.render(); - } finally { - // Clean up streaming resources regardless of outcome - this.cleanupStreamingResources(); } - } catch (error) { - // Use centralized error handler - ErrorHandler.handleError(error, 'ChatView.handleUserInput'); - // Update any streaming messages to non-streaming state to prevent stale messages - this.messages = this.messages.map((msg) => - msg.isStreaming ? { ...msg, isStreaming: false } : msg - ); - 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 { - if (this.sendButton) { - this.sendButton.disabled = false; - } + // Clean up streaming resources regardless of outcome + this.cleanupStreamingResources(); } } + + // State + private messages: ChatMessage[] = []; + private lastMessageEl: HTMLElement | null = null; + private newChatButton: HTMLElement | null = null; + private sendButton: HTMLElement | null = null; + private inputEl: HTMLTextAreaElement | null = null; + private chatContainer: HTMLElement | null = null; + private sendButtonClickHandler: (() => void) | null = null; + private inputKeyDownHandler: ((event: KeyboardEvent) => void) | null = null; + private newChatButtonClickHandler: (() => void) | null = null; + private sendButtonClickWrapper: (() => void) | null = null; + private inputKeyDownWrapper: ((event: KeyboardEvent) => void) | null = null; + private newChatButtonClickWrapper: (() => void) | null = null; + private listenersAttached: boolean = false; + private settings: PluginSettings; + 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[]; +} + +const MAX_STREAM_CHUNKS = 1000; diff --git a/src/constants.js b/src/constants.js index aa489ad..87fe449 100644 --- a/src/constants.js +++ b/src/constants.js @@ -1,18 +1,18 @@ -"use strict"; +'use strict'; // Default plugin settings -Object.defineProperty(exports, "__esModule", { value: true }); +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/constants.ts b/src/constants.ts index 50c2b36..132d55f 100644 --- a/src/constants.ts +++ b/src/constants.ts @@ -1,5 +1,3 @@ -// Default plugin settings - export const DEFAULT_SETTINGS = { ollamaUrl: 'http://localhost:11434', model: 'llama3', diff --git a/src/indexing-pipeline/extraction.js b/src/indexing-pipeline/extraction.js new file mode 100644 index 0000000..df48969 --- /dev/null +++ b/src/indexing-pipeline/extraction.js @@ -0,0 +1,88 @@ +"use strict"; +// src/indexing-pipeline/extraction.ts +Object.defineProperty(exports, "__esModule", { value: true }); +exports.ContentExtractor = void 0; +/** + * Extracts raw content from a vault file including: + * - Markdown content + * - YAML frontmatter + * - Headings + * - Embedded code blocks + * - First paragraph + */ +class ContentExtractor { + extractFromFile(file, content) { + const frontmatter = {}; + const headings = []; + const embeddedCodeBlocks = []; + let firstParagraph; + // Extract frontmatter + const frontmatterMatch = content.match(/^---(.*?)---/s); + 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; + } + } + else if (key.trim() === 'tags') { + if (value) { + frontmatter.tags = value; + } + } + else { + // Store other frontmatter fields as-is + frontmatter[key.trim()] = value; + } + } + } + catch { + // If frontmatter parsing fails, continue with empty frontmatter + } + } + // Extract headings + const headingMatches = content.match(/^#{1,6} (.*?)$/gm); + if (headingMatches) { + headings.push(...headingMatches.map((h) => h.replace(/^#{1,6} /, ''))); + } + // Extract embedded code blocks + const codeBlockMatches = content.match(/```([\s\S]*?)```/g); + if (codeBlockMatches) { + embeddedCodeBlocks.push(...codeBlockMatches); + } + // Extract first paragraph + const paragraphMatch = content.match(/^([^#]*?)(?=#|\n\n|$)/s); + if (paragraphMatch) { + firstParagraph = paragraphMatch[1].trim(); + } + return { + basename: file.basename, + path: file.path, + content, + frontmatter, + headings, + embeddedCodeBlocks, + firstParagraph, + }; + } + /** + * Extracts just the raw text content without headers, frontmatter, etc. + */ + extractRawText(content) { + return content + .replace(/^---.*?---/s, '') + .replace(/^#.*?$/gm, '') + .replace(/```.*?```/gs, '') + .replace(/`.*?`/g, '') + .replace(/\[.*?\]\(.*?\)/g, '') + .trim(); + } +} +exports.ContentExtractor = ContentExtractor; diff --git a/src/indexing-pipeline/extraction.ts b/src/indexing-pipeline/extraction.ts new file mode 100644 index 0000000..a332e72 --- /dev/null +++ b/src/indexing-pipeline/extraction.ts @@ -0,0 +1,109 @@ +// src/indexing-pipeline/extraction.ts + +// VaultFile interface is defined locally since it's not exported from types +interface VaultFile { + basename: string; + path: string; +} + +export interface Frontmatter { + title?: string; + tags?: string; + [key: string]: unknown; +} + +export interface ExtractedContent { + basename: string; + path: string; + content: string; + frontmatter: Frontmatter; + headings: string[]; + embeddedCodeBlocks: string[]; + firstParagraph?: string; +} + +/** + * Extracts raw content from a vault file including: + * - Markdown content + * - YAML frontmatter + * - Headings + * - Embedded code blocks + * - First paragraph + */ +export class ContentExtractor { + extractFromFile(file: VaultFile, content: string): ExtractedContent { + const frontmatter: Frontmatter = {}; + const headings: string[] = []; + const embeddedCodeBlocks: string[] = []; + let firstParagraph: string | undefined; + + // Extract frontmatter + const frontmatterMatch = content.match(/^---(.*?)---/s); + 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; + } + } else if (key.trim() === 'tags') { + if (value) { + frontmatter.tags = value; + } + } else { + // Store other frontmatter fields as-is + frontmatter[key.trim()] = value; + } + } + } catch { + // If frontmatter parsing fails, continue with empty frontmatter + } + } + + // Extract headings + const headingMatches = content.match(/^#{1,6} (.*?)$/gm); + if (headingMatches) { + headings.push(...headingMatches.map((h: string) => h.replace(/^#{1,6} /, ''))); + } + + // Extract embedded code blocks + const codeBlockMatches = content.match(/```([\s\S]*?)```/g); + if (codeBlockMatches) { + embeddedCodeBlocks.push(...codeBlockMatches); + } + + // Extract first paragraph + const paragraphMatch = content.match(/^([^#]*?)(?=#|\n\n|$)/s); + if (paragraphMatch) { + firstParagraph = paragraphMatch[1].trim(); + } + + return { + basename: file.basename, + path: file.path, + content, + frontmatter, + headings, + embeddedCodeBlocks, + firstParagraph, + }; + } + + /** + * Extracts just the raw text content without headers, frontmatter, etc. + */ + extractRawText(content: string): string { + return content + .replace(/^---.*?---/s, '') + .replace(/^#.*?$/gm, '') + .replace(/```.*?```/gs, '') + .replace(/`.*?`/g, '') + .replace(/\[.*?\]\(.*?\)/g, '') + .trim(); + } +} diff --git a/src/indexing-pipeline/index.js b/src/indexing-pipeline/index.js new file mode 100644 index 0000000..00eaba3 --- /dev/null +++ b/src/indexing-pipeline/index.js @@ -0,0 +1,12 @@ +"use strict"; +// src/indexing-pipeline/index.ts +Object.defineProperty(exports, "__esModule", { value: true }); +exports.IndexingPipeline = exports.ContentVectorizer = exports.ContentNormalizer = exports.ContentExtractor = void 0; +var extraction_1 = require("./extraction"); +Object.defineProperty(exports, "ContentExtractor", { enumerable: true, get: function () { return extraction_1.ContentExtractor; } }); +var normalization_1 = require("./normalization"); +Object.defineProperty(exports, "ContentNormalizer", { enumerable: true, get: function () { return normalization_1.ContentNormalizer; } }); +var vectorization_1 = require("./vectorization"); +Object.defineProperty(exports, "ContentVectorizer", { enumerable: true, get: function () { return vectorization_1.ContentVectorizer; } }); +var pipeline_1 = require("./pipeline"); +Object.defineProperty(exports, "IndexingPipeline", { enumerable: true, get: function () { return pipeline_1.IndexingPipeline; } }); diff --git a/src/indexing-pipeline/index.ts b/src/indexing-pipeline/index.ts new file mode 100644 index 0000000..ae1dc23 --- /dev/null +++ b/src/indexing-pipeline/index.ts @@ -0,0 +1,6 @@ +// src/indexing-pipeline/index.ts + +export { ContentExtractor } from './extraction'; +export { ContentNormalizer } from './normalization'; +export { ContentVectorizer } from './vectorization'; +export { IndexingPipeline } from './pipeline'; diff --git a/src/indexing-pipeline/normalization.js b/src/indexing-pipeline/normalization.js new file mode 100644 index 0000000..3fbf923 --- /dev/null +++ b/src/indexing-pipeline/normalization.js @@ -0,0 +1,159 @@ +"use strict"; +// src/indexing-pipeline/normalization.ts +Object.defineProperty(exports, "__esModule", { value: true }); +exports.ContentNormalizer = void 0; +/** + * Normalizes and enriches extracted content + */ +class ContentNormalizer { + /** + * Normalizes content by: + * - Standardizing dates to ISO 8601 + * - Converting to lowercase for tokenization + * - Extracting tokens + * - Adding metadata + */ + normalize(extractedContent) { + const { basename, path, content, frontmatter, headings, firstParagraph } = extractedContent; + // Standardize title (remove .md extension) + const title = basename.replace(/\.md$/, ''); + // Extract tokens (lowercase, remove stop words, etc.) + const tokens = this.tokenize(content); + // Normalize dates (if present in frontmatter) + const normalizedFrontmatter = this.normalizeFrontmatter(frontmatter); + // Calculate word count + const wordCount = content.split(/\s+/).filter(Boolean).length; + return { + path, + title, + content, + tokens, + headings, + frontmatter: normalizedFrontmatter, + firstParagraph, + wordCount, + // Add timestamps if available in frontmatter + createdAt: this.extractDate(frontmatter, 'created') || this.extractDate(frontmatter, 'date'), + updatedAt: this.extractDate(frontmatter, 'updated'), + }; + } + /** + * Tokenizes text content by splitting on whitespace and removing stop words + */ + tokenize(text) { + const stopWords = new Set([ + 'the', + 'a', + 'an', + 'and', + 'or', + 'but', + 'is', + 'are', + 'was', + 'were', + 'in', + 'on', + 'at', + 'to', + 'of', + 'for', + 'with', + 'as', + 'by', + 'it', + 'its', + 'that', + 'this', + 'these', + 'those', + 'from', + 'up', + 'out', + 'off', + 'over', + 'under', + 'again', + 'further', + 'then', + 'once', + 'here', + 'there', + 'when', + 'where', + 'why', + 'how', + 'all', + 'any', + 'both', + 'each', + 'few', + 'more', + 'most', + 'other', + 'some', + 'such', + 'no', + 'nor', + 'not', + 'only', + 'own', + 'same', + 'so', + 'than', + 'too', + 'very', + 'just', + 'now', + ]); + return text + .toLowerCase() + .split(/\W+/) + .filter((token) => token.length > 1 && !stopWords.has(token)); + } + /** + * Normalizes frontmatter by standardizing data types and formats + */ + normalizeFrontmatter(frontmatter) { + const normalized = {}; + for (const [key, value] of Object.entries(frontmatter)) { + if (key === 'tags' && typeof value === 'string') { + // Convert tag string to array if needed + normalized.tags = value.split(',').map((tag) => tag.trim()); + } + else if (key === 'date' || key === 'created' || key === 'updated') { + // Try to parse and standardize date formats + if (typeof value === 'string') { + const date = new Date(value); + if (!isNaN(date.getTime())) { + normalized[key] = date.toISOString(); + } + else { + normalized[key] = value; // Keep original if invalid date + } + } + else { + normalized[key] = value; + } + } + else { + normalized[key] = value; + } + } + return normalized; + } + /** + * Extracts a date from frontmatter + */ + extractDate(frontmatter, key) { + const value = frontmatter[key]; + if (typeof value === 'string') { + const date = new Date(value); + if (!isNaN(date.getTime())) { + return date.toISOString(); + } + } + return undefined; + } +} +exports.ContentNormalizer = ContentNormalizer; diff --git a/src/indexing-pipeline/normalization.ts b/src/indexing-pipeline/normalization.ts new file mode 100644 index 0000000..baa6044 --- /dev/null +++ b/src/indexing-pipeline/normalization.ts @@ -0,0 +1,207 @@ +// src/indexing-pipeline/normalization.ts + +// 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; + content: string; + tokens: string[]; + headings: string[]; + frontmatter: Record; + firstParagraph?: string; + // Additional enrichment fields + wordCount: number; + createdAt?: string; + updatedAt?: string; +} + +/** + * Normalizes and enriches extracted content + */ +export class ContentNormalizer { + /** + * Normalizes content by: + * - Standardizing dates to ISO 8601 + * - Converting to lowercase for tokenization + * - Extracting tokens + * - Adding metadata + */ + normalize(extractedContent: ExtractedContent): NormalizedContent { + const { basename, path, content, frontmatter, headings, firstParagraph } = extractedContent; + + // Standardize title (remove .md extension) + const title = basename.replace(/\.md$/, ''); + + // Extract tokens (lowercase, remove stop words, etc.) + const tokens = this.tokenize(content); + + // Normalize dates (if present in frontmatter) + const normalizedFrontmatter = this.normalizeFrontmatter(frontmatter); + + // Calculate word count + const wordCount = content.split(/\s+/).filter(Boolean).length; + + return { + path, + title, + content, + tokens, + headings, + frontmatter: normalizedFrontmatter, + firstParagraph, + wordCount, + // Add timestamps if available in frontmatter + createdAt: this.extractDate(frontmatter, 'created') || this.extractDate(frontmatter, 'date'), + updatedAt: this.extractDate(frontmatter, 'updated'), + }; + } + + /** + * Tokenizes text content by splitting on whitespace and removing stop words + */ + private tokenize(text: string): string[] { + const stopWords = new Set([ + 'the', + 'a', + 'an', + 'and', + 'or', + 'but', + 'is', + 'are', + 'was', + 'were', + 'in', + 'on', + 'at', + 'to', + 'of', + 'for', + 'with', + 'as', + 'by', + 'it', + 'its', + 'that', + 'this', + 'these', + 'those', + 'from', + 'up', + 'out', + 'off', + 'over', + 'under', + 'again', + 'further', + 'then', + 'once', + 'here', + 'there', + 'when', + 'where', + 'why', + 'how', + 'all', + 'any', + 'both', + 'each', + 'few', + 'more', + 'most', + 'other', + 'some', + 'such', + 'no', + 'nor', + 'not', + 'only', + 'own', + 'same', + 'so', + 'than', + 'too', + 'very', + 'just', + 'now', + ]); + + return text + .toLowerCase() + .split(/\W+/) + .filter((token) => token.length > 1 && !stopWords.has(token)); + } + + /** + * Normalizes frontmatter by standardizing data types and formats + */ + private normalizeFrontmatter(frontmatter: Record): Record { + const normalized: Record = {}; + + for (const [key, value] of Object.entries(frontmatter)) { + if (key === 'tags' && typeof value === 'string') { + // Convert tag string to array if needed + normalized.tags = value.split(',').map((tag) => tag.trim()); + } else if (key === 'date' || key === 'created' || key === 'updated') { + // Try to parse and standardize date formats + if (typeof value === 'string') { + const date = new Date(value); + if (!isNaN(date.getTime())) { + normalized[key] = date.toISOString(); + } else { + normalized[key] = value; // Keep original if invalid date + } + } else { + normalized[key] = value; + } + } else { + normalized[key] = value; + } + } + + return normalized; + } + + /** + * Extracts a date from frontmatter + */ + private extractDate(frontmatter: Record, key: string): string | undefined { + const value = frontmatter[key]; + if (typeof value === 'string') { + const date = new Date(value); + if (!isNaN(date.getTime())) { + return date.toISOString(); + } + } + return undefined; + } +} + +/** + * Interface for a normalized content chunk + */ +export interface ContentChunk { + id: string; + path: string; + title: string; + content: string; + tokens: string[]; + headings: string[]; + frontmatter: Record; + firstParagraph?: string; + wordCount: number; + createdAt?: string; + updatedAt?: string; + // Additional enrichment fields for vectorization + chunkIndex: number; + chunkSize: number; +} diff --git a/src/indexing-pipeline/pipeline.js b/src/indexing-pipeline/pipeline.js new file mode 100644 index 0000000..81243a0 --- /dev/null +++ b/src/indexing-pipeline/pipeline.js @@ -0,0 +1,64 @@ +"use strict"; +// src/indexing-pipeline/pipeline.ts +Object.defineProperty(exports, "__esModule", { value: true }); +exports.IndexingPipeline = void 0; +const extraction_1 = require("./extraction"); +const normalization_1 = require("./normalization"); +const vectorization_1 = require("./vectorization"); +class IndexingPipeline { + constructor(config) { + this.extractor = new extraction_1.ContentExtractor(); + this.normalizer = new normalization_1.ContentNormalizer(); + this.vectorizer = new vectorization_1.ContentVectorizer({ + model: config.embeddingModel, + ollamaUrl: config.ollamaUrl, + }); + } + /** + * Processes a vault file through the entire pipeline + */ + async processFile(file, content) { + try { + // Extraction step + const extracted = this.extractor.extractFromFile(file, content); + // Normalization/Enrichment step + const normalized = this.normalizer.normalize(extracted); + // Return the normalized content as an index entry + return { + path: normalized.path, + title: normalized.title, + content: this.extractor.extractRawText(content).substring(0, 500), + score: 0, // Score will be calculated during search + }; + } + catch (error) { + return null; + } + } + /** + * Processes multiple files in batches + */ + async 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 content = fileContents[file.path]; + if (!content) { + return null; + } + const entry = await 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); + } + return results; + } +} +exports.IndexingPipeline = IndexingPipeline; diff --git a/src/indexing-pipeline/pipeline.ts b/src/indexing-pipeline/pipeline.ts new file mode 100644 index 0000000..bd1760f --- /dev/null +++ b/src/indexing-pipeline/pipeline.ts @@ -0,0 +1,87 @@ +// src/indexing-pipeline/pipeline.ts + +import { VaultIndexEntry } from '../types'; +import { ContentExtractor } from './extraction'; +import { ContentNormalizer } from './normalization'; +import { ContentVectorizer } from './vectorization'; + +interface PipelineConfig { + ollamaUrl: string; + embeddingModel: string; +} + +export class IndexingPipeline { + private extractor: ContentExtractor; + private normalizer: ContentNormalizer; + private vectorizer: ContentVectorizer; + + constructor(config: PipelineConfig) { + this.extractor = new ContentExtractor(); + this.normalizer = new ContentNormalizer(); + this.vectorizer = new ContentVectorizer({ + model: config.embeddingModel, + ollamaUrl: config.ollamaUrl, + }); + } + + /** + * Processes a vault file through the entire pipeline + */ + async processFile(file: any, content: string): Promise { + try { + // Extraction step + const extracted = this.extractor.extractFromFile(file, content); + + // Normalization/Enrichment step + const normalized = this.normalizer.normalize(extracted); + + // Return the normalized content as an index entry + return { + path: normalized.path, + title: normalized.title, + content: this.extractor.extractRawText(content).substring(0, 500), + score: 0, // Score will be calculated during search + }; + } catch (error) { + return null; + } + } + + /** + * Processes multiple files in batches + */ + async processFilesInBatches( + files: any[], + fileContents: Record, + batchSize: number = 10 + ): Promise { + 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 content = fileContents[file.path]; + if (!content) { + return null; + } + + const entry = await 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 + ); + results.push(...validResults); + } + + return results; + } +} diff --git a/src/indexing-pipeline/vectorization.js b/src/indexing-pipeline/vectorization.js new file mode 100644 index 0000000..0e054eb --- /dev/null +++ b/src/indexing-pipeline/vectorization.js @@ -0,0 +1,55 @@ +"use strict"; +// src/indexing-pipeline/vectorization.ts +Object.defineProperty(exports, "__esModule", { value: true }); +exports.ContentVectorizer = void 0; +/** + * Vectorizes content chunks using Ollama embeddings + */ +class ContentVectorizer { + constructor(config, fetchFn) { + this.model = config.model; + this.ollamaUrl = config.ollamaUrl; + this.fetchFn = fetchFn ?? fetch; + } + /** + * Generates embeddings for a content chunk + */ + async vectorize(chunk) { + try { + const prompt = this.createPrompt(chunk); + const response = await this.fetchFn(`${this.ollamaUrl}/api/embeddings`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + model: this.model, + prompt: prompt, + }), + }); + if (!response.ok) { + throw new Error(`Embedding failed with status ${response.status}`); + } + const data = await response.json(); + return data.embedding; + } + catch (error) { + // Return empty array on failure to maintain compatibility + console.warn(`Failed to generate embedding: ${String(error)}`); + return []; + } + } + /** + * Creates a prompt from content chunk for embedding + */ + createPrompt(chunk) { + // Combine important elements for embedding + const parts = [ + chunk.title, + chunk.firstParagraph, + chunk.content.substring(0, 1000), // Limit content to avoid long prompts + chunk.headings.join(' '), + JSON.stringify(chunk.frontmatter), + ].filter(Boolean); + return parts.join('\n\n'); + } +} +exports.ContentVectorizer = ContentVectorizer; diff --git a/src/indexing-pipeline/vectorization.ts b/src/indexing-pipeline/vectorization.ts new file mode 100644 index 0000000..40fecc4 --- /dev/null +++ b/src/indexing-pipeline/vectorization.ts @@ -0,0 +1,68 @@ +// src/indexing-pipeline/vectorization.ts + +import { ContentChunk } from './normalization'; + +interface VectorizationConfig { + model: string; + ollamaUrl: string; +} + +/** + * Vectorizes content chunks using Ollama embeddings + */ +export class ContentVectorizer { + private model: string; + private ollamaUrl: string; + private fetchFn: typeof fetch; + + constructor(config: VectorizationConfig, fetchFn?: typeof fetch) { + this.model = config.model; + this.ollamaUrl = config.ollamaUrl; + this.fetchFn = fetchFn ?? fetch; + } + + /** + * Generates embeddings for a content chunk + */ + async vectorize(chunk: ContentChunk): Promise { + try { + const prompt = this.createPrompt(chunk); + + const response = await this.fetchFn(`${this.ollamaUrl}/api/embeddings`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + model: this.model, + prompt: prompt, + }), + }); + + if (!response.ok) { + throw new Error(`Embedding failed with status ${response.status}`); + } + + const data = await response.json(); + return data.embedding; + } catch (error) { + // Return empty array on failure to maintain compatibility + console.warn(`Failed to generate embedding: ${String(error)}`); + return []; + } + } + + /** + * Creates a prompt from content chunk for embedding + */ + private createPrompt(chunk: ContentChunk): string { + // Combine important elements for embedding + const parts = [ + chunk.title, + chunk.firstParagraph, + chunk.content.substring(0, 1000), // Limit content to avoid long prompts + chunk.headings.join(' '), + JSON.stringify(chunk.frontmatter), + ].filter(Boolean); + + return parts.join('\n\n'); + } +} diff --git a/src/main.js b/src/main.js index dd98b29..51342b6 100644 --- a/src/main.js +++ b/src/main.js @@ -1,199 +1,204 @@ -"use strict"; -Object.defineProperty(exports, "__esModule", { value: true }); -const obsidian_1 = require("obsidian"); -const chat_view_1 = require("./chat-view"); -const utils_1 = require("./utils"); -const constants_1 = require("./constants"); -const error_handler_1 = require("./error-handler"); +'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'; class OllamaPlugin extends obsidian_1.Plugin { - constructor() { - super(...arguments); - this.settings = constants_1.DEFAULT_SETTINGS; + 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.'); } - async onload() { - // Initialize logging - utils_1.Logger.info('Ollama Plugin loading...', 'plugin'); - await this.loadSettings(); - utils_1.Logger.info('Plugin loaded successfully', 'plugin'); - try { - this.registerView('ollama-chat-view', (leaf) => new chat_view_1.ChatView(leaf, this.settings)); - } - catch (error) { - utils_1.Logger.error('Failed to register view: ' + error.message, 'plugin'); - new obsidian_1.Notice('Failed to register Ollama chat view'); - // Don't throw - let the plugin continue loading other features - } - try { - this.addRibbonIcon('message-square', 'Ollama Chat', async () => { - const leaf = this.app.workspace.getLeaf(); - await leaf.setViewState({ - type: 'ollama-chat-view', - active: true, - }); - await this.app.workspace.revealLeaf(leaf); - }); - } - catch (error) { - utils_1.Logger.error('Failed to add ribbon icon: ' + error.message, 'plugin'); - new obsidian_1.Notice('Failed to add Ollama ribbon icon'); - // Don't throw - let the plugin continue loading other features - } - this.addSettingTab(new OllamaSettingTab(this.app, this)); + } + 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 loadSettings() { - try { - const data = (await this.loadData()); - if (data) { - utils_1.Logger.debug('Loading saved settings', 'settings'); - this.settings = { - ...constants_1.DEFAULT_SETTINGS, - ...data, - cacheConfig: { ...constants_1.DEFAULT_SETTINGS.cacheConfig, ...data.cacheConfig }, - }; - } - } - catch (error) { - error_handler_1.ErrorHandler.handleError(error, 'settings load'); - } - } - async saveSettings() { - try { - // Validate settings before saving - const validationErrors = (0, utils_1.validatePluginSettings)(this.settings); - if (validationErrors.length > 0) { - utils_1.Logger.error('Validation errors prevented saving settings: ' + validationErrors.join('; '), 'settings'); - new obsidian_1.Notice(`Cannot save settings: ${validationErrors[0]}`); - return false; - } - utils_1.Logger.debug('Saving settings: ' + JSON.stringify(this.settings), 'settings'); - await this.saveData(this.settings); - utils_1.Logger.info('Settings saved successfully', 'settings'); - return true; - } - catch (error) { - error_handler_1.ErrorHandler.handleError(error, 'settings save'); - return false; - } - } - // Notify all open ChatView instances when settings change - notifyChatViews() { - const leaves = this.app.workspace.getLeavesOfType('ollama-chat-view'); - leaves.forEach((leaf) => { - const view = leaf.view; - if (view instanceof chat_view_1.ChatView) { - view.onSettingsChange(this.settings); - } - }); - } - async clearSemanticCache() { - const leaves = this.app.workspace.getLeavesOfType('ollama-chat-view'); - for (const leaf of leaves) { - const view = leaf.view; - if (view instanceof chat_view_1.ChatView) { - await view.clearCache(); - return; - } - } + } + async clearSemanticCache() { + if (this.semanticCache) { + await this.semanticCache.clear(); } + } + 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.default = OllamaPlugin; +exports.OllamaPlugin = OllamaPlugin; class OllamaSettingTab extends obsidian_1.PluginSettingTab { - constructor(app, plugin) { - super(app, plugin); - this.plugin = plugin; - } - display() { - // Clear any existing content first to prevent duplicates - this.containerEl.empty(); - // Create container for settings - const container = this.containerEl.createDiv(); - new obsidian_1.Setting(container) - .setName('Ollama URL') - .setDesc('URL of your Ollama instance') - .addText((text) => text.setValue(this.plugin.settings.ollamaUrl).onChange(async (value) => { - const urlValidation = (0, utils_1.validateOllamaUrl)(value); - if (urlValidation.valid) { - utils_1.Logger.debug('URL changed to: ' + value, 'settings'); - this.plugin.settings.ollamaUrl = value; - await this.plugin.saveSettings(); - this.plugin.notifyChatViews(); - } - else { - utils_1.Logger.warn('Invalid URL format: ' + urlValidation.error, 'settings'); - new obsidian_1.Notice(urlValidation.error || 'Invalid Ollama URL format.'); - } - })); - new obsidian_1.Setting(container) - .setName('Model') - .setDesc('Model to use for chat') - .addText((text) => text.setValue(this.plugin.settings.model).onChange(async (value) => { - const modelValidation = (0, utils_1.validateModelName)(value); - if (modelValidation.valid) { - utils_1.Logger.debug('Model changed to: ' + value, 'settings'); - this.plugin.settings.model = value; - await this.plugin.saveSettings(); - this.plugin.notifyChatViews(); - } - else { - utils_1.Logger.warn('Invalid model name format: ' + modelValidation.error, 'settings'); - new obsidian_1.Notice(modelValidation.error || 'Invalid model name format.'); - } - })); - new obsidian_1.Setting(container) - .setName('Enable Semantic Cache') - .setDesc('Cache responses semantically to speed up repeated queries') - .addToggle((toggle) => toggle.setValue(this.plugin.settings.cacheConfig.enabled).onChange(async (value) => { - this.plugin.settings.cacheConfig.enabled = value; + 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; await this.plugin.saveSettings(); - this.plugin.notifyChatViews(); - })); - new obsidian_1.Setting(container) - .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) => { + } 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(container) - .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(container) - .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(); + 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.'); } - else { - new obsidian_1.Notice('Similarity threshold must be a number between 0 and 1.'); - } - })); - new obsidian_1.Setting(container) - .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(); - } + }) + ); + 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 bbd708c..78cb641 100755 --- a/src/main.ts +++ b/src/main.ts @@ -1,165 +1,159 @@ import { Plugin, WorkspaceLeaf, App, Setting, Notice, PluginSettingTab } from 'obsidian'; import { ChatView } from './chat-view'; -import { PluginSettings } from './types'; -import { validatePluginSettings, validateOllamaUrl, validateModelName, Logger } from './utils'; import { DEFAULT_SETTINGS } from './constants'; -import { ErrorHandler } from './error-handler'; +import { SemanticCache } from './semantic-cache'; export default class OllamaPlugin extends Plugin { - settings: PluginSettings = DEFAULT_SETTINGS; + settings = DEFAULT_SETTINGS; + semanticCache?: SemanticCache; async onload() { - // Initialize logging - Logger.info('Ollama Plugin loading...', 'plugin'); - await this.loadSettings(); - Logger.info('Plugin loaded successfully', 'plugin'); - try { - this.registerView( - 'ollama-chat-view', - (leaf: WorkspaceLeaf) => new ChatView(leaf, this.settings) - ); - } catch (error) { - Logger.error('Failed to register view: ' + (error as Error).message, 'plugin'); - new Notice('Failed to register Ollama chat view'); - // Don't throw - let the plugin continue loading other features - } + // Register the chat view + this.registerView( + 'ollama-chat-view', + (leaf: WorkspaceLeaf) => new ChatView(leaf, this.settings) + ); - try { - this.addRibbonIcon('message-square', 'Ollama Chat', async () => { - const leaf = this.app.workspace.getLeaf(); - await leaf.setViewState({ - type: 'ollama-chat-view', - active: true, - }); - await this.app.workspace.revealLeaf(leaf); - }); - } catch (error) { - Logger.error('Failed to add ribbon icon: ' + (error as Error).message, 'plugin'); - new Notice('Failed to add Ollama ribbon icon'); - // Don't throw - let the plugin continue loading other features - } + // 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 Notice('Semantic cache cleared.'); + }, + }); + + // Add a settings tab this.addSettingTab(new OllamaSettingTab(this.app, this)); + + // Initialize the semantic cache + this.semanticCache = new SemanticCache(this.settings.cacheConfig); + try { + await this.semanticCache.initialize(); + } catch (error) { + console.error('Failed to initialize semantic cache:', error); + new Notice('Semantic cache initialization failed. Check console for details.'); + } + } + + async onunload() { + this.unregisterView('ollama-chat-view'); } async loadSettings() { - try { - const data = (await this.loadData()) as Partial | null; - if (data) { - Logger.debug('Loading saved settings', 'settings'); - this.settings = { - ...DEFAULT_SETTINGS, - ...data, - cacheConfig: { ...DEFAULT_SETTINGS.cacheConfig, ...data.cacheConfig }, - }; - } - } catch (error) { - ErrorHandler.handleError(error, 'settings load'); - } + this.settings = Object.assign({}, DEFAULT_SETTINGS, await this.loadData()); } async saveSettings() { - try { - // Validate settings before saving - const validationErrors = validatePluginSettings(this.settings); - if (validationErrors.length > 0) { - Logger.error( - 'Validation errors prevented saving settings: ' + validationErrors.join('; '), - 'settings' - ); - new Notice(`Cannot save settings: ${validationErrors[0]}`); - return false; - } + await this.saveData(this.settings); + } - Logger.debug('Saving settings: ' + JSON.stringify(this.settings), 'settings'); - await this.saveData(this.settings); - Logger.info('Settings saved successfully', 'settings'); - return true; - } catch (error) { - ErrorHandler.handleError(error, 'settings save'); - return false; + async activateChatView() { + const existing = this.app.workspace.getLeavesOfType('ollama-chat-view'); + if (existing.length > 0) { + this.app.workspace.revealLeaf(existing[0]); + } else { + await this.app.workspace.getRightLeaf(false).setViewState({ + type: 'ollama-chat-view', + active: true, + }); } } - // Notify all open ChatView instances when settings change - public notifyChatViews(): void { + async clearSemanticCache() { + if (this.semanticCache) { + await this.semanticCache.clear(); + } + } + + notifyChatViews() { const leaves = this.app.workspace.getLeavesOfType('ollama-chat-view'); leaves.forEach((leaf) => { - const view = leaf.view; - if (view instanceof ChatView) { - view.onSettingsChange(this.settings); + if (leaf.view instanceof ChatView) { + leaf.view.updateSettings(this.settings); } }); } - - public async clearSemanticCache(): Promise { - const leaves = this.app.workspace.getLeavesOfType('ollama-chat-view'); - for (const leaf of leaves) { - const view = leaf.view; - if (view instanceof ChatView) { - await view.clearCache(); - return; - } - } - } } class OllamaSettingTab extends PluginSettingTab { - private plugin: OllamaPlugin; + plugin: OllamaPlugin; constructor(app: App, plugin: OllamaPlugin) { super(app, plugin); this.plugin = plugin; } - display(): void { - // Clear any existing content first to prevent duplicates - this.containerEl.empty(); + display() { + const { containerEl } = this; + containerEl.empty(); + containerEl.createEl('h2', { text: 'Ollama Settings' }); - // Create container for settings - const container = this.containerEl.createDiv() as HTMLElement; - - new Setting(container) + new Setting(containerEl) .setName('Ollama URL') - .setDesc('URL of your Ollama instance') + .setDesc('URL for your Ollama instance (default: http://localhost:11434)') .addText((text) => text.setValue(this.plugin.settings.ollamaUrl).onChange(async (value) => { - const urlValidation = validateOllamaUrl(value); - if (urlValidation.valid) { - Logger.debug('URL changed to: ' + value, 'settings'); - this.plugin.settings.ollamaUrl = value; - await this.plugin.saveSettings(); - this.plugin.notifyChatViews(); - } else { - Logger.warn('Invalid URL format: ' + urlValidation.error, 'settings'); - new Notice(urlValidation.error || 'Invalid Ollama URL format.'); - } + this.plugin.settings.ollamaUrl = value; + await this.plugin.saveSettings(); }) ); - new Setting(container) + new Setting(containerEl) .setName('Model') - .setDesc('Model to use for chat') + .setDesc('Ollama model to use (default: llama3)') .addText((text) => text.setValue(this.plugin.settings.model).onChange(async (value) => { - const modelValidation = validateModelName(value); - if (modelValidation.valid) { - Logger.debug('Model changed to: ' + value, 'settings'); - this.plugin.settings.model = value; + this.plugin.settings.model = value; + await this.plugin.saveSettings(); + }) + ); + + new 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(); - this.plugin.notifyChatViews(); } else { - Logger.warn('Invalid model name format: ' + modelValidation.error, 'settings'); - new Notice(modelValidation.error || 'Invalid model name format.'); + new Notice('Vault search limit must be a positive integer.'); } }) ); - new Setting(container) + new 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 Notice('Max message history must be a positive integer.'); + } + }) + ); + + new Setting(containerEl) .setName('Enable Semantic Cache') - .setDesc('Cache responses semantically to speed up repeated queries') + .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; @@ -168,7 +162,7 @@ class OllamaSettingTab extends PluginSettingTab { }) ); - new Setting(container) + new Setting(containerEl) .setName('ChromaDB URL') .setDesc('URL for your ChromaDB instance (default: http://localhost:8000)') .addText((text) => @@ -180,7 +174,7 @@ class OllamaSettingTab extends PluginSettingTab { }) ); - new Setting(container) + new Setting(containerEl) .setName('Cache Embedding Model') .setDesc('Ollama model used to generate embeddings for the semantic cache') .addText((text) => @@ -191,7 +185,7 @@ class OllamaSettingTab extends PluginSettingTab { }) ); - new Setting(container) + new Setting(containerEl) .setName('Cache Similarity Threshold') .setDesc( 'Minimum cosine similarity (0–1) for a cache hit. Higher values require closer matches.' @@ -210,7 +204,7 @@ class OllamaSettingTab extends PluginSettingTab { }) ); - new Setting(container) + new Setting(containerEl) .setName('Clear Semantic Cache') .setDesc('Delete all cached responses from ChromaDB') .addButton((button) => @@ -225,7 +219,7 @@ class OllamaSettingTab extends PluginSettingTab { ); } - hide(): void { + hide() { // Clear the container to prevent duplicate elements this.containerEl.empty(); } diff --git a/src/ollama-client.js b/src/ollama-client.js index f073cdb..af6bb27 100644 --- a/src/ollama-client.js +++ b/src/ollama-client.js @@ -1,299 +1,179 @@ -"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); - } + 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 initializeCache() { + if (this.cacheService) { + await this.cacheService.initialize(); } - async clearCache() { - if (this.cacheService) { - await this.cacheService.clearCache(); - } + } + async clearCache() { + if (this.cacheService) { + await this.cacheService.clearCache(); } - cancelStream() { - if (this.currentStreamController) { - this.currentStreamController.abort(); - this.currentStreamController = null; - } + } + 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 - let lastUserMsg = messages[messages.length - 1]; - if (lastUserMsg && lastUserMsg.role !== 'user') { - // Find the last user message if not the last one - for (let i = messages.length - 1; i >= 0; i--) { - if (messages[i].role === 'user') { - lastUserMsg = messages[i]; - break; - } - } - } - if (lastUserMsg && this.cacheService) { - const cached = await this.cacheService.getCache(lastUserMsg.content); - if (cached) { - yield { role: 'assistant', content: cached, tool_calls: [] }; - return; - } - } - if (this.cacheService && lastUserMsg) { - 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(''); - void this.cacheService.setCache(lastUserMsg.content, fullContent); - } - else { - yield* this.streamChatWithRetry(messages, tools, 0); - } + } + 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; } - async streamChatAsPromise(messages, tools = []) { - const chunks = []; - for await (const chunk of this.streamChat(messages, tools)) { - chunks.push(chunk); - } - return chunks; + // 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; + } } - async chat(messages, tools = []) { - // Bypass cache if tools are involved - if (tools.length > 0) { - return this.chatWithRetry(messages, tools, 0); - } - // Find the last user message - let lastUserMsg = messages[messages.length - 1]; - if (lastUserMsg && lastUserMsg.role !== 'user') { - // Find the last user message if not the last one - for (let i = messages.length - 1; i >= 0; i--) { - if (messages[i].role === 'user') { - lastUserMsg = messages[i]; - break; - } - } - } - if (lastUserMsg && this.cacheService) { - const cached = await this.cacheService.getCache(lastUserMsg.content); - if (cached) { - return { role: 'assistant', content: cached, tool_calls: [] }; - } - } - const response = await this.chatWithRetry(messages, tools, 0); - if (this.cacheService && lastUserMsg) { - void this.cacheService.setCache(lastUserMsg.content, response.content); - } - return response; + const chunks = []; + for await (const chunk of this.streamChatWithRetry(messages, tools, 0)) { + chunks.push(chunk); + yield chunk; } - async *streamChatWithRetry(messages, tools = [], attempt = 0) { - const controller = new AbortController(); - this.currentStreamController = controller; + 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; + } + 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); + } + } + } + if (buffer.trim() !== '') { try { - const response = await this.fetchFn(`${this.baseURL}/api/chat`, { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - }, - body: JSON.stringify({ - model: this.model, - messages, - tools, - stream: true, - }), - signal: controller.signal, - }); - if (!response) { - throw new Error('No response received from Ollama API'); - } - if (!response.ok) { - if (response.status >= 500 && attempt < this.maxRetries) { - const retryDelay = Math.pow(2, attempt) * 100; - utils_1.Logger.warn(`Network error (status ${response.status}), retrying in ${retryDelay}ms (attempt ${attempt + 1}/${this.maxRetries})`, 'ollama-client'); - // Wait for the delay, but also check if the stream was cancelled during backoff - await new Promise((resolve, reject) => { - const timer = setTimeout(() => resolve(), retryDelay); - // If the controller was aborted during the delay, cancel the retry - if (controller.signal.aborted) { - clearTimeout(timer); - const abortError = controller.signal.reason ?? - new DOMException('The operation was aborted.', 'AbortError'); - reject(abortError); - return; - } - controller.signal.addEventListener('abort', () => { - clearTimeout(timer); - const abortError = controller.signal.reason ?? - new DOMException('The operation was aborted.', 'AbortError'); - reject(abortError); - }, { once: true }); - }); - // Only proceed with retry if this controller is still the active one. - // If a newer stream replaced currentStreamController during backoff, - // abandon the retry to avoid overwriting the newer stream's controller. - if (this.currentStreamController !== controller) { - return; - } - yield* this.streamChatWithRetry(messages, tools, attempt + 1); - // After successful retry, we need to return (not continue processing this response) - return; - } - else { - 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('ndjson') && !contentType.includes('json'))) { - throw new Error('Invalid response format'); - } - const reader = response.body.getReader(); - const decoder = new TextDecoder(); - let buffer = ''; - let malformedCount = 0; - const maxMalformed = 50; - try { - while (true) { - const { done, value } = await reader.read(); - if (done) - break; - buffer += decoder.decode(value, { stream: true }); - const lines = buffer.split('\n'); - buffer = lines.pop() ?? ''; - for (const line of lines) { - if (!line.trim()) - continue; - try { - const parsed = JSON.parse(line); - this.throwIfOllamaError(parsed); - const message = this.toOllamaMessage(parsed.message); - if (!message) { - continue; - } - malformedCount = 0; - yield message; - } - catch (error) { - if (error instanceof Error && error.message.startsWith('Ollama error:')) { - throw error; - } - malformedCount++; - if (malformedCount > maxMalformed) { - throw new Error('Too many malformed chunks in stream'); - } - utils_1.Logger.warn(`Skipped malformed chunk: ${line.substring(0, 80)}... - ${error.message}`, 'ollama-client'); - } - } - } - if (buffer.trim()) { - try { - const parsed = JSON.parse(buffer); - this.throwIfOllamaError(parsed); - const message = this.toOllamaMessage(parsed.message); - if (message) { - yield message; - } - } - catch (error) { - if (error instanceof Error && error.message.startsWith('Ollama error:')) { - throw error; - } - utils_1.Logger.warn(`Failed to parse final chunk: ${buffer.substring(0, 80)}...`, 'ollama-client'); - } - } - } - finally { - reader.releaseLock(); - } - } - catch (error) { - if (!(error instanceof Error) || error.name !== 'AbortError') { - utils_1.Logger.error(`Stream encountered an error: ${String(error)}`, 'ollama-client'); - throw error; - } - // Re-throw the abort error to allow stream consumers to handle it - throw error; - } - finally { - // Only clear controller if it's still the current one (not replaced by a new stream) - if (this.currentStreamController === controller) { - controller.abort(); - this.currentStreamController = null; - } + const parsed = JSON.parse(buffer); + yield parsed.message; + } catch (error) { + console.error('Failed to parse final chunk:', buffer, error); } + } + } 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 = [], attempt = 0) { - const controller = new AbortController(); - try { - const response = await this.fetchFn(`${this.baseURL}/api/chat`, { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - }, - body: JSON.stringify({ - model: this.model, - messages, - tools, - stream: false, - }), - signal: controller.signal, - }); - if (!response) { - throw new Error('No response received from Ollama API'); - } - if (!response.ok) { - if (response.status >= 500 && attempt < this.maxRetries) { - const retryDelay = Math.pow(2, attempt) * 100; - utils_1.Logger.warn(`Network error (status ${response.status}), retrying in ${retryDelay}ms (attempt ${attempt + 1}/${this.maxRetries})`, 'ollama-client'); - await new Promise((resolve) => setTimeout(resolve, retryDelay)); - return this.chatWithRetry(messages, tools, attempt + 1); - } - throw new types_1.ApiError(`Ollama API error: ${response.status}`, response.status); - } - const data = (await response.json()); - return (this.toOllamaMessage(data.message) ?? { role: 'assistant', content: '', tool_calls: [] }); - } - finally { - controller.abort(); - } - } - throwIfOllamaError(parsed) { - if (parsed.error) { - const errorMsg = typeof parsed.error === 'string' ? parsed.error : JSON.stringify(parsed.error); - throw new Error(`Ollama error: ${errorMsg}`); - } - } - toOllamaMessage(value) { - if (!value || typeof value !== 'object') { - return null; - } - const record = value; - const toolCalls = Array.isArray(record.tool_calls) ? record.tool_calls : []; - return { - role: record.role ?? 'assistant', - content: typeof record.content === 'string' ? record.content : '', - tool_calls: toolCalls, - }; + } + 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; + } } + } } exports.OllamaClient = OllamaClient; diff --git a/src/ollama-client.ts b/src/ollama-client.ts index c02091f..5d60322 100644 --- a/src/ollama-client.ts +++ b/src/ollama-client.ts @@ -1,12 +1,13 @@ // src/ollama-client.ts import type { OllamaMessage, OllamaTool } from './types'; -import { ApiError, CacheConfig } from './types'; +import { ApiError } from './types'; import { Logger } from './utils'; import { SemanticCacheService } from './semantic-cache'; interface OllamaChatResponse { message?: Partial; + error?: string; } export class OllamaClient { @@ -17,7 +18,7 @@ export class OllamaClient { private currentStreamController: AbortController | null = null; private cacheService?: SemanticCacheService; - constructor(baseURL: string, model: string, fetchFn?: typeof fetch, cacheConfig?: CacheConfig) { + constructor(baseURL: string, model: string, fetchFn?: typeof fetch, cacheConfig?: any) { this.baseURL = baseURL; this.model = model; this.fetchFn = fetchFn ?? fetch; @@ -57,17 +58,7 @@ export class OllamaClient { } // Find the last user message - let lastUserMsg: OllamaMessage | undefined = messages[messages.length - 1]; - if (lastUserMsg && lastUserMsg.role !== 'user') { - // Find the last user message if not the last one - for (let i = messages.length - 1; i >= 0; i--) { - if (messages[i].role === 'user') { - lastUserMsg = messages[i]; - break; - } - } - } - + const lastUserMsg = [...messages].reverse().find((m) => m.role === 'user'); if (lastUserMsg && this.cacheService) { const cached = await this.cacheService.getCache(lastUserMsg.content); if (cached) { @@ -76,68 +67,44 @@ export class OllamaClient { } } - if (this.cacheService && lastUserMsg) { - const chunks: OllamaMessage[] = []; - for await (const chunk of this.streamChatWithRetry(messages, tools, 0)) { - chunks.push(chunk); - yield chunk; - } - const fullContent = chunks.map((c) => c.content).join(''); - void this.cacheService.setCache(lastUserMsg.content, fullContent); - } else { - yield* this.streamChatWithRetry(messages, tools, 0); - } - } - - async streamChatAsPromise( - messages: OllamaMessage[], - tools: OllamaTool[] = [] - ): Promise { const chunks: OllamaMessage[] = []; - for await (const chunk of this.streamChat(messages, tools)) { + 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); } - return chunks; } async chat(messages: OllamaMessage[], tools: OllamaTool[] = []): Promise { - // Bypass cache if tools are involved + // 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 - let lastUserMsg: OllamaMessage | undefined = messages[messages.length - 1]; - if (lastUserMsg && lastUserMsg.role !== 'user') { - // Find the last user message if not the last one - for (let i = messages.length - 1; i >= 0; i--) { - if (messages[i].role === 'user') { - lastUserMsg = messages[i]; - break; - } - } - } - + 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, tool_calls: [] }; + 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; } - private async *streamChatWithRetry( + async *streamChatWithRetry( messages: OllamaMessage[], tools: OllamaTool[] = [], - attempt: number = 0 + retryCount: number ): AsyncGenerator { const controller = new AbortController(); this.currentStreamController = controller; @@ -150,171 +117,76 @@ export class OllamaClient { }, body: JSON.stringify({ model: this.model, - messages, - tools, + messages: messages, stream: true, + tools: tools, }), signal: controller.signal, }); - if (!response) { - throw new Error('No response received from Ollama API'); - } - if (!response.ok) { - if (response.status >= 500 && attempt < this.maxRetries) { - const retryDelay = Math.pow(2, attempt) * 100; - Logger.warn( - `Network error (status ${response.status}), retrying in ${retryDelay}ms (attempt ${attempt + 1}/${this.maxRetries})`, - 'ollama-client' - ); - - // Wait for the delay, but also check if the stream was cancelled during backoff - await new Promise((resolve, reject) => { - const timer = setTimeout(() => resolve(), retryDelay); - - // If the controller was aborted during the delay, cancel the retry - if (controller.signal.aborted) { - clearTimeout(timer); - const abortError: Error = - controller.signal.reason ?? - (new DOMException('The operation was aborted.', 'AbortError') as unknown as Error); - reject(abortError); - return; - } - - controller.signal.addEventListener( - 'abort', - () => { - clearTimeout(timer); - const abortError: Error = - controller.signal.reason ?? - (new DOMException( - 'The operation was aborted.', - 'AbortError' - ) as unknown as Error); - reject(abortError); - }, - { once: true } - ); - }); - - // Only proceed with retry if this controller is still the active one. - // If a newer stream replaced currentStreamController during backoff, - // abandon the retry to avoid overwriting the newer stream's controller. - if (this.currentStreamController !== controller) { - return; - } - - yield* this.streamChatWithRetry(messages, tools, attempt + 1); - // After successful retry, we need to return (not continue processing this response) - return; - } else { - throw new ApiError(`Ollama API error: ${response.status}`, response.status); - } + throw new Error(`HTTP error! status: ${response.status}`); } if (!response.body) { - throw new Error('No response body'); - } - - const contentType = response.headers.get('content-type'); - if (!contentType || (!contentType.includes('ndjson') && !contentType.includes('json'))) { - throw new Error('Invalid response format'); + throw new Error('Response body is null'); } const reader = response.body.getReader(); const decoder = new TextDecoder(); let buffer = ''; - let malformedCount = 0; - const maxMalformed = 50; - try { - while (true) { - const { done, value } = await reader.read(); - if (done) break; - - buffer += decoder.decode(value, { stream: true }); - - const lines = buffer.split('\n'); - buffer = lines.pop() ?? ''; - - for (const line of lines) { - if (!line.trim()) continue; - - try { - const parsed = JSON.parse(line) as Record; - this.throwIfOllamaError(parsed); - - const message = this.toOllamaMessage(parsed.message); - if (!message) { - continue; - } - - malformedCount = 0; - yield message; - } catch (error) { - if (error instanceof Error && error.message.startsWith('Ollama error:')) { - throw error; - } - - malformedCount++; - if (malformedCount > maxMalformed) { - throw new Error('Too many malformed chunks in stream'); - } - - Logger.warn( - `Skipped malformed chunk: ${line.substring(0, 80)}... - ${(error as Error).message}`, - 'ollama-client' - ); - } - } + while (true) { + const { done, value } = await reader.read(); + if (done) { + break; } - if (buffer.trim()) { + 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(buffer) as Record; - this.throwIfOllamaError(parsed); - - const message = this.toOllamaMessage(parsed.message); - if (message) { - yield message; - } + const parsed = JSON.parse(line); + yield parsed.message; } catch (error) { - if (error instanceof Error && error.message.startsWith('Ollama error:')) { - throw error; - } - Logger.warn( - `Failed to parse final chunk: ${buffer.substring(0, 80)}...`, - 'ollama-client' - ); + // Log but don't throw - malformed JSON is not critical + console.error('Failed to parse chunk:', line, error); } } - } finally { - reader.releaseLock(); + } + + if (buffer.trim() !== '') { + try { + const parsed = JSON.parse(buffer); + yield parsed.message; + } catch (error) { + console.error('Failed to parse final chunk:', buffer, error); + } } } catch (error) { - if (!(error instanceof Error) || error.name !== 'AbortError') { - Logger.error(`Stream encountered an error: ${String(error)}`, 'ollama-client'); + 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; } - // Re-throw the abort error to allow stream consumers to handle it - throw error; } finally { - // Only clear controller if it's still the current one (not replaced by a new stream) - if (this.currentStreamController === controller) { - controller.abort(); - this.currentStreamController = null; - } + this.currentStreamController = null; } } - private async chatWithRetry( + async chatWithRetry( messages: OllamaMessage[], tools: OllamaTool[] = [], - attempt: number = 0 + retryCount: number ): Promise { - const controller = new AbortController(); try { const response = await this.fetchFn(`${this.baseURL}/api/chat`, { method: 'POST', @@ -323,59 +195,29 @@ export class OllamaClient { }, body: JSON.stringify({ model: this.model, - messages, - tools, - stream: false, + messages: messages, + tools: tools, }), - signal: controller.signal, }); - if (!response) { - throw new Error('No response received from Ollama API'); - } - if (!response.ok) { - if (response.status >= 500 && attempt < this.maxRetries) { - const retryDelay = Math.pow(2, attempt) * 100; - Logger.warn( - `Network error (status ${response.status}), retrying in ${retryDelay}ms (attempt ${attempt + 1}/${this.maxRetries})`, - 'ollama-client' - ); - await new Promise((resolve) => setTimeout(resolve, retryDelay)); - return this.chatWithRetry(messages, tools, attempt + 1); - } - throw new ApiError(`Ollama API error: ${response.status}`, response.status); + 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()) as OllamaChatResponse; - return ( - this.toOllamaMessage(data.message) ?? { role: 'assistant', content: '', tool_calls: [] } - ); - } finally { - controller.abort(); + 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; + } } } - - private throwIfOllamaError(parsed: Record): void { - if (parsed.error) { - const errorMsg = - typeof parsed.error === 'string' ? parsed.error : JSON.stringify(parsed.error); - throw new Error(`Ollama error: ${errorMsg}`); - } - } - - private toOllamaMessage(value: unknown): OllamaMessage | null { - if (!value || typeof value !== 'object') { - return null; - } - - const record = value as Partial; - const toolCalls = Array.isArray(record.tool_calls) ? record.tool_calls : []; - - return { - role: record.role ?? 'assistant', - content: typeof record.content === 'string' ? record.content : '', - tool_calls: toolCalls, - }; - } } diff --git a/src/semantic-cache.js b/src/semantic-cache.js index 303afbf..17187f8 100644 --- a/src/semantic-cache.js +++ b/src/semantic-cache.js @@ -1,118 +1,96 @@ -"use strict"; +'use strict'; // 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 }); + 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; } - 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: ${String(error)}`, '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) { + utils_1.Logger.warn(`Cache lookup failed: ${error.message}`, 'semantic-cache'); + return null; } - async getEmbedding(text) { - try { - 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(`Embedding failed with status ${response.status}`); - } - const data = await response.json(); - return data.embedding; - } - catch (error) { - utils_1.Logger.warn(`Failed to generate embedding: ${String(error)}`, 'semantic-cache'); - return []; - } + } + 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(prompt) { - if (!this.collection || !this.config.enabled || !prompt.trim()) { - return null; - } - try { - const embedding = await this.getEmbedding(prompt); - if (!embedding.length) - return null; - const results = await this.collection.query({ - queryEmbeddings: [embedding], - nResults: 1, - include: ['metadatas', 'distances'], - }); - // Cosine distance = 1 - cosine_similarity - // We want distance < (1 - threshold) - if (results.distances && - results.distances[0] && - results.distances[0][0] < 1 - this.config.similarityThreshold) { - utils_1.Logger.debug('Semantic cache hit', 'semantic-cache'); - return results.metadatas?.[0]?.[0]?.fullResponse ?? null; - } - } - catch (error) { - utils_1.Logger.warn(`Cache lookup failed: ${String(error)}`, '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(prompt, response) { - if (!this.collection || !this.config.enabled || !prompt.trim() || !response.trim()) { - return; - } - try { - const embedding = await this.getEmbedding(prompt); - if (!embedding.length) - return; - // Fallback for crypto.randomUUID() if not available - let id; - if (typeof crypto !== 'undefined' && typeof crypto.randomUUID === 'function') { - id = crypto.randomUUID(); - } - else { - // Fallback to a simple ID generator if crypto is not available - id = 'cache_' + Date.now() + '_' + Math.random().toString(36).substr(2, 9); - } - const collectionAddResult = await this.collection.add({ - ids: [id], - embeddings: [embedding], - metadatas: [{ fullResponse: response }], - }); - utils_1.Logger.debug('Cached new response', 'semantic-cache'); - } - catch (error) { - utils_1.Logger.warn(`Cache write failed: ${String(error)}`, 'semantic-cache'); - } - } - async clearCache() { - if (this.collection && this.config.enabled) { - 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: ${String(error)}`, '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; + } } exports.SemanticCacheService = SemanticCacheService; diff --git a/src/semantic-cache.ts b/src/semantic-cache.ts index a52ff6c..f85ceca 100644 --- a/src/semantic-cache.ts +++ b/src/semantic-cache.ts @@ -6,20 +6,19 @@ import { CacheConfig } from './types'; export class SemanticCacheService { private client: ChromaClient; - private collection: any | null = null; + private collection: ReturnType | null = null; private config: CacheConfig; private ollamaURL: string; - private chromaURL: string; constructor(ollamaURL: string, config: CacheConfig) { 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 ChromaClient({ path: this.chromaURL }); + const chromaURL = config.chromaURL || 'http://localhost:8000'; + this.client = new ChromaClient({ path: chromaURL }); } - async initialize() { + async initialize(): Promise { if (!this.config.enabled) return; try { @@ -27,104 +26,83 @@ export class SemanticCacheService { name: this.config.collectionName, metadata: { 'hnsw:space': 'cosine' }, }); + Logger.info(`Semantic cache initialized: ${this.config.collectionName}`, 'semantic-cache'); } catch (error) { - Logger.error(`Failed to initialize semantic cache: ${String(error)}`, 'semantic-cache'); + Logger.error(`Failed to initialize semantic cache: ${error.message}`, 'semantic-cache'); + throw error; } } - private async getEmbedding(text: string): Promise { + async getCache(query: string): Promise { + if (!this.config.enabled || !this.collection) return null; + try { - const response = await fetch(`${this.ollamaURL}/api/embeddings`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ - model: this.config.embeddingModel, - prompt: text, - }), + const results = await this.collection.query({ + query_embeddings: await this.generateEmbedding(query), + n_results: 1, + where: { source: 'ollama' }, }); - if (!response.ok) { - throw new Error(`Embedding failed with status ${response.status}`); + 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; + } } - const data = await response.json(); - return data.embedding; + return null; } catch (error) { - Logger.warn(`Failed to generate embedding: ${String(error)}`, 'semantic-cache'); - return []; - } - } - - async getCache(prompt: string): Promise { - if (!this.collection || !this.config.enabled || !prompt.trim()) { + Logger.warn(`Cache lookup failed: ${error.message}`, 'semantic-cache'); return null; } - - try { - const embedding = await this.getEmbedding(prompt); - if (!embedding.length) return null; - - const results: any = await this.collection.query({ - queryEmbeddings: [embedding], - nResults: 1, - include: ['metadatas', 'distances'], - }); - - // Cosine distance = 1 - cosine_similarity - // We want distance < (1 - threshold) - if ( - results.distances && - results.distances[0] && - results.distances[0][0] < 1 - this.config.similarityThreshold - ) { - Logger.debug('Semantic cache hit', 'semantic-cache'); - return results.metadatas?.[0]?.[0]?.fullResponse ?? null; - } - } catch (error) { - Logger.warn(`Cache lookup failed: ${String(error)}`, 'semantic-cache'); - } - - return null; } - async setCache(prompt: string, response: string): Promise { - if (!this.collection || !this.config.enabled || !prompt.trim() || !response.trim()) { - return; - } + async setCache(query: string, response: string): Promise { + if (!this.config.enabled || !this.collection) return; try { - const embedding = await this.getEmbedding(prompt); - if (!embedding.length) return; - - // Fallback for crypto.randomUUID() if not available - let id: string; - if (typeof crypto !== 'undefined' && typeof crypto.randomUUID === 'function') { - id = crypto.randomUUID(); - } else { - // Fallback to a simple ID generator if crypto is not available - id = 'cache_' + Date.now() + '_' + Math.random().toString(36).substr(2, 9); - } - - const collectionAddResult: any = await this.collection.add({ - ids: [id], - embeddings: [embedding], - metadatas: [{ fullResponse: response }], + await this.collection.add({ + ids: [crypto.randomUUID()], + documents: [response], + embeddings: await this.generateEmbedding(query), + metadatas: [{ source: 'ollama' }], }); - Logger.debug('Cached new response', 'semantic-cache'); } catch (error) { - Logger.warn(`Cache write failed: ${String(error)}`, 'semantic-cache'); + Logger.warn(`Cache set failed: ${error.message}`, 'semantic-cache'); } } async clearCache(): Promise { - if (this.collection && this.config.enabled) { - try { - await this.collection.reset(); - Logger.info('Semantic cache cleared', 'semantic-cache'); - } catch (error) { - Logger.error(`Failed to clear semantic cache: ${String(error)}`, 'semantic-cache'); - } + if (!this.config.enabled || !this.collection) return; + + try { + await this.collection.reset(); + Logger.info('Semantic cache cleared', 'semantic-cache'); + } catch (error) { + Logger.error(`Failed to clear semantic cache: ${error.message}`, 'semantic-cache'); } } + + private async generateEmbedding(text: string): Promise { + 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; + } } diff --git a/src/types.ts b/src/types.ts index 26d1db8..d1bd0e3 100644 --- a/src/types.ts +++ b/src/types.ts @@ -30,60 +30,71 @@ export class NetworkError extends OllamaError { constructor(message: string, statusCode?: number) { super(message, ErrorType.NETWORK_ERROR); this.statusCode = statusCode; - Object.setPrototypeOf(this, NetworkError.prototype); } } export class ApiError extends OllamaError { - public readonly statusCode?: number; + public readonly statusCode: number; - constructor(message: string, statusCode?: number) { + constructor(message: string, statusCode: number) { super(message, ErrorType.API_ERROR); this.statusCode = statusCode; - Object.setPrototypeOf(this, ApiError.prototype); } } -export interface ValidationFieldDetails { - field?: string; - message?: string; -} - export class ValidationError extends OllamaError { - public readonly details?: ValidationFieldDetails; - - constructor(message: string, details?: ValidationFieldDetails) { + constructor(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 { - public readonly toolName: string; - - constructor(message: string, toolName: string) { + constructor(message: string) { super(message, ErrorType.TOOL_EXECUTION_ERROR); - this.toolName = toolName; - Object.setPrototypeOf(this, ToolExecutionError.prototype); } } -export class PathValidationError extends OllamaError { - public readonly path: string; +// ============================================================ +// Ollama Types +// ============================================================ - constructor(message: string, path: string) { - super(message, ErrorType.PATH_VALIDATION_ERROR); - this.path = path; - Object.setPrototypeOf(this, PathValidationError.prototype); - } +export interface OllamaMessage { + role: 'system' | 'user' | 'assistant'; + content: string; + tool_calls?: OllamaToolCall[]; +} + +export interface OllamaToolCall { + id: string; + type: 'function'; + function: { + name: string; + arguments: string; + }; +} + +export interface OllamaTool { + type: 'function'; + function: { + name: string; + description: string; + parameters: { + type: 'object'; + properties: { + [key: string]: { + type: string; + description?: string; + }; + }; + required?: string[]; + }; + }; } // ============================================================ @@ -106,82 +117,3 @@ export interface PluginSettings { lastIndexTime: number; cacheConfig: CacheConfig; } - -// ============================================================ -// Ollama Protocol Types -// ============================================================ - -export interface OllamaTool { - type: 'function'; - function: { - name: string; - description: string; - parameters: { - type: 'object'; - properties: Record; - required: string[]; - }; - }; -} - -export interface OllamaToolCall { - id: string; - type: 'function'; - function: { - name: string; - arguments: string | Record; - }; -} - -export interface OllamaMessage { - role: 'system' | 'user' | 'assistant' | 'tool'; - content: string; - tool_calls?: OllamaToolCall[]; -} - -// ============================================================ -// Tool Execution Types -// ============================================================ - -export interface ToolCall { - id: string; - type: 'function'; - function: { - name: string; - arguments: string | Record; - }; -} - -export interface ToolResult { - success: boolean; - message: string; -} - -export interface ExecutionResult { - success: boolean; - output: string; -} - -// ============================================================ -// Chat Message Types -// ============================================================ - -export interface ChatMessage { - id: string; - role: 'user' | 'assistant'; - content: string; - timestamp: number; - isStreaming?: boolean; - tool_calls?: ToolCall[]; -} - -// ============================================================ -// Vault Index Types -// ============================================================ - -export interface VaultIndexEntry { - path: string; - title: string; - content: string; - score: number; -} diff --git a/src/vault-indexer.js b/src/vault-indexer.js index d840c1e..103fb8d 100644 --- a/src/vault-indexer.js +++ b/src/vault-indexer.js @@ -1,268 +1,151 @@ -"use strict"; +'use strict'; // src/vault-indexer.ts -Object.defineProperty(exports, "__esModule", { value: true }); -exports.InMemoryCache = exports.VaultIndexer = void 0; -exports.createVaultIndexerWithCache = createVaultIndexerWithCache; -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(); - } - clear() { - this.store.clear(); - return Promise.resolve(); - } -} -exports.InMemoryCache = InMemoryCache; +Object.defineProperty(exports, '__esModule', { value: true }); +exports.VaultIndexer = void 0; +const obsidian_1 = require('obsidian'); +const utils_1 = require('./utils'); class VaultIndexer { - constructor(vault, cache) { - this.vault = null; - // Define weights for scoring - this.SCORING_WEIGHTS = { - HEADING: 5, - FRONTMATTER_TITLE: 3, - FRONTMATTER_TAGS: 2.5, - FIRST_PARAGRAPH: 1.5, - TOKEN: 1, - }; - this.vault = vault; - this.cache = cache; + 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); + 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'); + } } - async searchVault(query, limit = 5) { - if (!query || !query.trim()) { - return []; + 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 (!this.vault) { - throw new Error('Vault-like object not provided to VaultIndexer'); + if (entry.frontmatter?.tags && this.exactMatch(entry.frontmatter.tags, queryToken)) { + score += this.SCORING_WEIGHTS.FRONTMATTER_TAGS; + matched = true; } - const cacheKey = `query:${query.trim()}:limit:${limit}`; - if (this.cache) { - let cachedResults; - try { - cachedResults = await this.cache.get(cacheKey); - } - catch { - // Ignore cache retrieval errors and continue with normal processing - cachedResults = null; - } - if (cachedResults) { - try { - const parsedResults = JSON.parse(cachedResults); - return parsedResults.slice(0, limit); - } - catch { - // Ignore cache parse errors and continue with normal processing - } - } + if (entry.headings.some((heading) => heading.toLowerCase().includes(stemmed))) { + score += this.SCORING_WEIGHTS.HEADINGS; + matched = true; } - const queryTokens = this.tokenize(query.trim()); - const vault = this.vault; - const allFiles = vault.getMarkdownFiles(); - const results = await this.processFilesInBatches(vault, allFiles, queryTokens); - const filteredResults = results.sort((a, b) => b.score - a.score).slice(0, limit); - if (this.cache) { - try { - await this.cache.put(cacheKey, JSON.stringify(filteredResults)); - } - catch (error) { - utils_1.Logger.warn(`Failed to cache results for query "${query}": ${error instanceof Error ? error.message : String(error)}`, 'vault-indexer'); - } + if (entry.content.toLowerCase().includes(stemmed)) { + score += this.SCORING_WEIGHTS.CONTENT; + matched = true; } - return filteredResults; + 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; + } + } else if (key.trim() === 'tags') { + if (value) { + frontmatter.tags = value; + } + } + } + } catch { + utils_1.Logger.warn('Failed to parse frontmatter', 'vault-indexer'); + } } - async processFilesInBatches(vault, files, queryTokens) { - const 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) => { - try { - const content = await vault.read(file); - const tokenized = this.tokenizeContent(content); - const scoreResult = this.calculateWeightedScore(tokenized, queryTokens, file); - if (scoreResult.score > 0) { - const entry = { - path: file.path, - title: file.basename.replace(/\.md$/, ''), - content: content.substring(0, 500), - score: scoreResult.score, - }; - if (!seenPaths.has(entry.path)) { - seenPaths.add(entry.path); - return entry; - } - return null; - } - return null; - } - catch (error) { - utils_1.Logger.warn(`Failed to process ${file.path}: ${error instanceof Error ? error.message : String(error)}`, 'vault-indexer'); - return null; - } - })); - const validResults = batchResults.filter((result) => result !== null); - results.push(...validResults); - } - return results; - } - tokenize(text) { - const stopWords = new Set([ - 'the', - 'a', - 'an', - 'and', - 'or', - 'but', - 'is', - 'are', - 'was', - 'were', - 'in', - 'on', - 'at', - 'to', - 'of', - 'for', - 'with', - 'as', - 'by', - 'it', - 'its', - 'that', - 'this', - 'these', - 'those', - ]); - return text - .toLowerCase() - .split(/\W+/) - .filter((token) => token.length > 1 && !stopWords.has(token)); - } - tokenizeContent(content) { - const tokens = []; - const headings = []; - const frontmatter = {}; - let firstParagraph; - const frontmatterMatch = content.match(/^---(.*?)---/s); - 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; - } - } - else if (key.trim() === 'tags') { - if (value) { - frontmatter.tags = value; - } - } - } - } - catch { - utils_1.Logger.warn('Failed to parse frontmatter', 'vault-indexer'); - } - } - const headingMatches = content.match(/^#{1,6} (.*?)$/gm); - if (headingMatches) { - headings.push(...headingMatches.map((h) => h.replace(/^#{1,6} /, ''))); - } - const paragraphMatch = content.match(/^([^#]*?)(?=#|\n\n|$)/s); - if (paragraphMatch) { - firstParagraph = paragraphMatch[1].trim(); - } - const allText = content - .replace(/^---.*?---/s, '') - .replace(/^#.*?$/gm, '') - .replace(/```.*?```/gs, '') - .replace(/`.*?`/g, '') - .replace(/\[.*?\]\(.*?\)/g, ''); - tokens.push(...this.tokenize(allText)); - return { tokens, headings, frontmatter, firstParagraph }; - } - calculateWeightedScore(tokenized, queryTokens, file) { - let totalScore = 0; - const matchedTokens = new Set(); - for (const queryToken of queryTokens) { - let tokenScore = 0; - const stemmed = this.stemToken(queryToken); - let matched = false; - if (tokenized.frontmatter?.title && - this.exactMatch(tokenized.frontmatter.title, queryToken)) { - tokenScore += this.SCORING_WEIGHTS.FRONTMATTER_TITLE; - matched = true; - } - else if (file && - file.basename && - this.exactMatch(file.basename.replace(/\.md$/, ''), queryToken)) { - tokenScore += this.SCORING_WEIGHTS.FRONTMATTER_TITLE; - matched = true; - } - if (tokenized.frontmatter?.tags && this.exactMatch(tokenized.frontmatter.tags, queryToken)) { - tokenScore += this.SCORING_WEIGHTS.FRONTMATTER_TAGS; - matched = true; - } - if (tokenized.headings.some((heading) => heading.toLowerCase().includes(stemmed))) { - tokenScore += this.SCORING_WEIGHTS.HEADING; - matched = true; - } - if (tokenized.firstParagraph && tokenized.firstParagraph.toLowerCase().includes(stemmed)) { - tokenScore += this.SCORING_WEIGHTS.FIRST_PARAGRAPH; - matched = true; - } - if (tokenized.tokens.includes(stemmed)) { - tokenScore += this.SCORING_WEIGHTS.TOKEN; - matched = true; - } - if (matched) { - totalScore += tokenScore; - matchedTokens.add(queryToken); - } - } - return { - score: totalScore, - matchedFields: Array.from(matchedTokens), - }; - } - stemToken(token) { - // Improved stemmer that handles edge cases - // - // Limitations: - // - Simple suffix removal (e.g., 's', 'ed', 'ing') may lead to over-stemming or incorrect roots. - // - For example, stemming "mice" results in "mic", which is incorrect. - // - Consider using a more robust NLP library if the plugin environment permits. - // - if (token.length <= 3) - return token; // Don't stem very short tokens - if (token.endsWith('s')) - return token.slice(0, -1); - if (token.endsWith('ed') && token.length > 4) - return token.slice(0, -2); // Don't stem 3-letter words ending in ed - if (token.endsWith('ing') && token.length > 5) - return token.slice(0, -3); // Don't stem 4-letter words ending in ing - return token; - } - exactMatch(content, token) { - const stemmedToken = this.stemToken(token); - return content.toLowerCase().includes(stemmedToken); + 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 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; -// Convenience method to create a VaultIndexer with an in-memory cache -function createVaultIndexerWithCache(vault) { - return new VaultIndexer(vault, new InMemoryCache()); -} diff --git a/src/vault-indexer.ts b/src/vault-indexer.ts index bedf4a2..8ce3d7c 100644 --- a/src/vault-indexer.ts +++ b/src/vault-indexer.ts @@ -1,74 +1,18 @@ // src/vault-indexer.ts -import { VaultIndexEntry } from './types'; +import { VaultLike } from 'obsidian'; import { Logger } from './utils'; +import { Cache } from './cache'; -interface Cache { - get(key: string): Promise; - put(key: string, value: string): Promise; - clear(): Promise; -} - -class InMemoryCache implements Cache { - private store: Map; - - constructor() { - this.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(); - } - - clear(): Promise { - this.store.clear(); - return Promise.resolve(); - } -} - -interface Frontmatter { - title?: string; - tags?: string; -} - -interface VaultFile { - basename: string; - path: string; -} - -interface VaultLike { - getMarkdownFiles(): VaultFile[]; - read(file: VaultFile): Promise; -} - -interface TokenizedContent { - tokens: string[]; - headings: string[]; - frontmatter: Frontmatter; - firstParagraph?: string; -} - -interface ScoreResult { - score: number; - matchedFields: string[]; -} - -class VaultIndexer { +export class VaultIndexer { private vault: VaultLike | null = null; private cache?: Cache; - - // Define weights for scoring private readonly SCORING_WEIGHTS = { - HEADING: 5, - FRONTMATTER_TITLE: 3, - FRONTMATTER_TAGS: 2.5, - FIRST_PARAGRAPH: 1.5, - TOKEN: 1, + TITLE: 5, + FRONTMATTER_TITLE: 4, + FRONTMATTER_TAGS: 3, + HEADINGS: 2, + CONTENT: 1, }; constructor(vault: VaultLike, cache?: Cache) { @@ -76,15 +20,33 @@ class VaultIndexer { this.cache = cache; } - async searchVault(query: string, limit: number = 5): Promise { + async getVaultEntries() { + const files = this.vault.getMarkdownFiles(); + const entries = []; + for (const file of files) { + try { + const content = await this.vault.cachedRead(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) { + Logger.warn(`Failed to read file ${file.path}: ${error.message}`, 'vault-indexer'); + } + } + return entries; + } + + async searchVault(query: string, limit = 3) { if (!query || !query.trim()) { return []; } - if (!this.vault) { - throw new Error('Vault-like object not provided to VaultIndexer'); - } - const cacheKey = `query:${query.trim()}:limit:${limit}`; if (this.cache) { let cachedResults; @@ -96,7 +58,7 @@ class VaultIndexer { } if (cachedResults) { try { - const parsedResults = JSON.parse(cachedResults) as VaultIndexEntry[]; + const parsedResults = JSON.parse(cachedResults); return parsedResults.slice(0, limit); } catch { // Ignore cache parse errors and continue with normal processing @@ -104,118 +66,89 @@ class VaultIndexer { } } - const queryTokens = this.tokenize(query.trim()); - const vault = this.vault; - const allFiles = vault.getMarkdownFiles(); - const results = await this.processFilesInBatches(vault, allFiles, queryTokens); + 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 filteredResults = results.sort((a, b) => b.score - a.score).slice(0, limit); + // 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); if (this.cache) { try { - await this.cache.put(cacheKey, JSON.stringify(filteredResults)); + await this.cache.put(cacheKey, JSON.stringify(results)); } catch (error) { Logger.warn( - `Failed to cache results for query "${query}": ${error instanceof Error ? error.message : String(error)}`, + `Failed to cache results for query "${query}": ${error.message}`, 'vault-indexer' ); } } - return filteredResults; - } - - private async processFilesInBatches( - vault: VaultLike, - files: VaultFile[], - queryTokens: string[] - ): Promise { - const batchSize = 10; - 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) => { - try { - const content = await vault.read(file); - const tokenized = this.tokenizeContent(content); - const scoreResult = this.calculateWeightedScore(tokenized, queryTokens, file); - if (scoreResult.score > 0) { - const entry: VaultIndexEntry = { - path: file.path, - title: file.basename.replace(/\.md$/, ''), - content: content.substring(0, 500), - score: scoreResult.score, - }; - if (!seenPaths.has(entry.path)) { - seenPaths.add(entry.path); - return entry; - } - return null; - } - return null; - } catch (error) { - Logger.warn( - `Failed to process ${file.path}: ${error instanceof Error ? error.message : String(error)}`, - 'vault-indexer' - ); - return null; - } - }) - ); - - const validResults = batchResults.filter( - (result): result is NonNullable => result !== null - ); - results.push(...validResults); - } - return results; } - private tokenize(text: string): string[] { - const stopWords = new Set([ - 'the', - 'a', - 'an', - 'and', - 'or', - 'but', - 'is', - 'are', - 'was', - 'were', - 'in', - 'on', - 'at', - 'to', - 'of', - 'for', - 'with', - 'as', - 'by', - 'it', - 'its', - 'that', - 'this', - 'these', - 'those', - ]); - return text + private tokenizeQuery(query: string) { + const tokens = query .toLowerCase() - .split(/\W+/) - .filter((token) => token.length > 1 && !stopWords.has(token)); + .replace(/[^\w\s]/g, '') + .split(/\s+/) + .filter((token) => token.length > 0); + return { + tokens, + stemmedTokens: tokens.map((token) => this.stemToken(token)), + }; } - private tokenizeContent(content: string): TokenizedContent { - const tokens: string[] = []; - const headings: string[] = []; - const frontmatter: Frontmatter = {}; - let firstParagraph: string | undefined; + private stemToken(token: string) { + // Simple stemming for now - in a real implementation, this would be more sophisticated + return token; + } - const frontmatterMatch = content.match(/^---(.*?)---/s); + private exactMatch(text: string | undefined, queryToken: string) { + if (!text) return false; + return text.toLowerCase().includes(queryToken.toLowerCase()); + } + + private parseMarkdown(content: string) { + const frontmatterRegex = /^---\n([\s\S]*?)\n---\n/; + const frontmatterMatch = content.match(frontmatterRegex); + const frontmatter = {}; if (frontmatterMatch) { try { const frontmatterContent = frontmatterMatch[1]; @@ -239,111 +172,30 @@ class VaultIndexer { } } - const headingMatches = content.match(/^#{1,6} (.*?)$/gm); - if (headingMatches) { - headings.push(...headingMatches.map((h: string) => h.replace(/^#{1,6} /, ''))); + 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 paragraphMatch = content.match(/^([^#]*?)(?=#|\n\n|$)/s); - if (paragraphMatch) { - firstParagraph = paragraphMatch[1].trim(); - } - - const allText = content - .replace(/^---.*?---/s, '') - .replace(/^#.*?$/gm, '') - .replace(/```.*?```/gs, '') - .replace(/`.*?`/g, '') - .replace(/\[.*?\]\(.*?\)/g, ''); - tokens.push(...this.tokenize(allText)); - - return { tokens, headings, frontmatter, firstParagraph }; - } - - private calculateWeightedScore( - tokenized: TokenizedContent, - queryTokens: string[], - file?: VaultFile - ): ScoreResult { - let totalScore = 0; - const matchedTokens: Set = new Set(); - - for (const queryToken of queryTokens) { - let tokenScore = 0; - const stemmed = this.stemToken(queryToken); - let matched = false; - - if ( - tokenized.frontmatter?.title && - this.exactMatch(tokenized.frontmatter.title, queryToken) - ) { - tokenScore += this.SCORING_WEIGHTS.FRONTMATTER_TITLE; - matched = true; - } else if ( - file && - file.basename && - this.exactMatch(file.basename.replace(/\.md$/, ''), queryToken) - ) { - tokenScore += this.SCORING_WEIGHTS.FRONTMATTER_TITLE; - matched = true; - } - - if (tokenized.frontmatter?.tags && this.exactMatch(tokenized.frontmatter.tags, queryToken)) { - tokenScore += this.SCORING_WEIGHTS.FRONTMATTER_TAGS; - matched = true; - } - - if (tokenized.headings.some((heading: string) => heading.toLowerCase().includes(stemmed))) { - tokenScore += this.SCORING_WEIGHTS.HEADING; - matched = true; - } - - if (tokenized.firstParagraph && tokenized.firstParagraph.toLowerCase().includes(stemmed)) { - tokenScore += this.SCORING_WEIGHTS.FIRST_PARAGRAPH; - matched = true; - } - - if (tokenized.tokens.includes(stemmed)) { - tokenScore += this.SCORING_WEIGHTS.TOKEN; - matched = true; - } - - if (matched) { - totalScore += tokenScore; - matchedTokens.add(queryToken); - } - } + 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 { - score: totalScore, - matchedFields: Array.from(matchedTokens), + frontmatter, + title, + headings, + content: contentWithoutFrontmatterAndHeadings, }; } - - private stemToken(token: string): string { - // Improved stemmer that handles edge cases - // - // Limitations: - // - Simple suffix removal (e.g., 's', 'ed', 'ing') may lead to over-stemming or incorrect roots. - // - For example, stemming "mice" results in "mic", which is incorrect. - // - Consider using a more robust NLP library if the plugin environment permits. - // - if (token.length <= 3) return token; // Don't stem very short tokens - if (token.endsWith('s')) return token.slice(0, -1); - if (token.endsWith('ed') && token.length > 4) return token.slice(0, -2); // Don't stem 3-letter words ending in ed - if (token.endsWith('ing') && token.length > 5) return token.slice(0, -3); // Don't stem 4-letter words ending in ing - return token; - } - - private exactMatch(content: string, token: string): boolean { - const stemmedToken = this.stemToken(token); - return content.toLowerCase().includes(stemmedToken); - } -} - -export { VaultIndexer, Cache, InMemoryCache }; - -// Convenience method to create a VaultIndexer with an in-memory cache -export function createVaultIndexerWithCache(vault: VaultLike): VaultIndexer { - return new VaultIndexer(vault, new InMemoryCache()); } diff --git a/tests/indexing-pipeline.test.ts b/tests/indexing-pipeline.test.ts new file mode 100644 index 0000000..205f796 --- /dev/null +++ b/tests/indexing-pipeline.test.ts @@ -0,0 +1,308 @@ +import { ContentExtractor } from '../src/indexing-pipeline/extraction'; +import { ContentNormalizer } from '../src/indexing-pipeline/normalization'; +import { ContentVectorizer } from '../src/indexing-pipeline/vectorization'; +import { IndexingPipeline } from '../src/indexing-pipeline/pipeline'; + +// Mock VaultFile interface for testing +interface MockVaultFile { + basename: string; + path: string; +} + +describe('Indexing Pipeline Components', () => { + describe('ContentExtractor', () => { + let extractor: ContentExtractor; + + beforeEach(() => { + extractor = new ContentExtractor(); + }); + + it('should extract frontmatter correctly', () => { + const content = `--- +title: Test Title +tags: algorithm, programming +date: 2023-01-01 +--- + +# Heading + +Content here`; + + const file: MockVaultFile = { basename: 'test', path: 'test.md' }; + const extracted = extractor.extractFromFile(file, content); + + expect(extracted.frontmatter.title).toBe('Test Title'); + expect(extracted.frontmatter.tags).toBe('algorithm, programming'); + expect(extracted.frontmatter.date).toBe('2023-01-01'); + expect(extracted.headings).toContain('Heading'); + }); + + it('should extract headings correctly', () => { + const content = `# Heading 1 +## Heading 2 +### Heading 3 + +Content`; + + const file: MockVaultFile = { basename: 'test', path: 'test.md' }; + const extracted = extractor.extractFromFile(file, content); + + expect(extracted.headings).toEqual(['Heading 1', 'Heading 2', 'Heading 3']); + }); + + it('should extract embedded code blocks', () => { + const content = `# Code Example + +\`\`\`javascript +console.log('hello world'); +\`\`\` + +Some content`; + + const file: MockVaultFile = { basename: 'test', path: 'test.md' }; + const extracted = extractor.extractFromFile(file, content); + + expect(extracted.embeddedCodeBlocks).toHaveLength(1); + expect(extracted.embeddedCodeBlocks[0]).toContain('console.log'); + }); + + it('should extract first paragraph', () => { + const content = `First paragraph here. + +Second paragraph here. + +# Heading`; + + const file: MockVaultFile = { basename: 'test', path: 'test.md' }; + const extracted = extractor.extractFromFile(file, content); + + expect(extracted.firstParagraph).toBe('First paragraph here.'); + }); + + it('should extract raw text correctly', () => { + const content = `--- +title: Test +--- + +# Heading + +Content with **bold** and [link](url). + +\`\`\`javascript +code +\`\`\``; + + const rawText = extractor.extractRawText(content); + expect(rawText).not.toContain('---'); + expect(rawText).not.toContain('# Heading'); + expect(rawText).not.toContain('```javascript'); + expect(rawText).toContain('Content with bold and link'); + }); + }); + + describe('ContentNormalizer', () => { + let normalizer: ContentNormalizer; + let extractor: ContentExtractor; + + beforeEach(() => { + normalizer = new ContentNormalizer(); + extractor = new ContentExtractor(); + }); + + it('should normalize frontmatter dates to ISO format', () => { + const content = `--- +title: Test +date: 2023-01-01 +created: 2023-06-15 +updated: invalid-date +tags: algorithm +--- + +Content`; + + const file: MockVaultFile = { basename: 'test', path: 'test.md' }; + const extracted = extractor.extractFromFile(file, content); + const normalized = normalizer.normalize(extracted); + + expect(normalized.frontmatter.date).toBe('2023-01-01T00:00:00.000Z'); + expect(normalized.frontmatter.created).toBe('2023-06-15T00:00:00.000Z'); + expect(normalized.frontmatter.updated).toBe('invalid-date'); // Should preserve invalid dates + }); + + it('should convert tags to array format', () => { + const content = `--- +title: Test +tags: algorithm, programming, javascript +--- + +Content`; + + const file: MockVaultFile = { basename: 'test', path: 'test.md' }; + const extracted = extractor.extractFromFile(file, content); + const normalized = normalizer.normalize(extracted); + + expect(normalized.frontmatter.tags).toEqual(['algorithm', 'programming', 'javascript']); + }); + + it('should calculate word count correctly', () => { + const content = `# Title + +This is a test document with several words to count. +It has multiple sentences and words to make it longer.`; + + const file: MockVaultFile = { basename: 'test', path: 'test.md' }; + const extracted = extractor.extractFromFile(file, content); + const normalized = normalizer.normalize(extracted); + + expect(normalized.wordCount).toBeGreaterThan(0); + }); + + it('should extract tokens correctly', () => { + const content = `# Test Document + +This is a test document with important keywords.`; + + const file: MockVaultFile = { basename: 'test', path: 'test.md' }; + const extracted = extractor.extractFromFile(file, content); + const normalized = normalizer.normalize(extracted); + + expect(normalized.tokens).toContain('test'); + expect(normalized.tokens).toContain('document'); + expect(normalized.tokens).toContain('important'); + expect(normalized.tokens).toContain('keywords'); + }); + + it('should extract title correctly', () => { + const content = `# Test Document + +Content`; + + const file: MockVaultFile = { basename: 'test.md', path: 'test.md' }; + const extracted = extractor.extractFromFile(file, content); + const normalized = normalizer.normalize(extracted); + + expect(normalized.title).toBe('test'); + }); + }); + + describe('ContentVectorizer', () => { + let vectorizer: ContentVectorizer; + + beforeEach(() => { + vectorizer = new ContentVectorizer({ + model: 'nomic-embed-text', + ollamaUrl: 'http://localhost:11434', + }); + }); + + it('should create a proper prompt from content chunk', () => { + const mockChunk = { + id: 'test', + path: 'test.md', + title: 'Test', + content: 'Test content', + tokens: ['test', 'content'], + headings: ['Heading'], + frontmatter: { tags: ['test'] }, + firstParagraph: 'First paragraph', + wordCount: 2, + chunkIndex: 0, + chunkSize: 100 + }; + + const prompt = (vectorizer as any).createPrompt(mockChunk); + + expect(prompt).toContain('Test'); + expect(prompt).toContain('First paragraph'); + expect(prompt).toContain('Heading'); + expect(prompt).toContain('tags'); + }); + + // Note: Actual embedding tests would require mocking fetch or integration testing + it('should handle vectorization errors gracefully', async () => { + // This test would require mocking fetch to simulate error responses + // For now, we're just ensuring the method exists and doesn't crash + const mockChunk = { + id: 'test', + path: 'test.md', + title: 'Test', + content: 'Test content', + tokens: ['test', 'content'], + headings: ['Heading'], + frontmatter: { tags: ['test'] }, + firstParagraph: 'First paragraph', + wordCount: 2, + chunkIndex: 0, + chunkSize: 100 + }; + + // Mock fetch to simulate an error + const originalFetch = global.fetch; + (global.fetch as any) = jest.fn().mockRejectedValue(new Error('Network error')); + + try { + const result = await vectorizer.vectorize(mockChunk); + expect(result).toEqual([]); + } finally { + global.fetch = originalFetch; + } + }); + }); + + describe('IndexingPipeline', () => { + let pipeline: IndexingPipeline; + + beforeEach(() => { + pipeline = new IndexingPipeline({ + ollamaUrl: 'http://localhost:11434', + embeddingModel: 'nomic-embed-text', + }); + }); + + it('should process files through the pipeline', async () => { + const file: MockVaultFile = { basename: 'test', path: 'test.md' }; + const content = `--- +title: Test Document +tags: test, example +--- + +# Introduction + +This is a test document for pipeline processing.`; + + const result = await pipeline.processFile(file, content); + + expect(result).not.toBeNull(); + expect(result?.title).toBe('test'); + expect(result?.path).toBe('test.md'); + expect(result?.content).toContain('This is a test document for pipeline processing'); + }); + + it('should handle processing errors gracefully', async () => { + const file: MockVaultFile = { basename: 'test', path: 'test.md' }; + + const result = await pipeline.processFile(file, ''); + + // Should not crash, but might return null or incomplete result + expect(result).toBeNull(); // Empty content should return null + }); + + it('should process files in batches', async () => { + const files: MockVaultFile[] = [ + { basename: 'file1', path: 'file1.md' }, + { basename: 'file2', path: 'file2.md' } + ]; + + const fileContents = { + 'file1.md': '# File 1\n\nContent 1', + 'file2.md': '# File 2\n\nContent 2' + }; + + const results = await pipeline.processFilesInBatches(files, fileContents, 1); + + expect(results).toHaveLength(2); + expect(results[0].title).toBe('file1'); + expect(results[1].title).toBe('file2'); + }); + }); +}); diff --git a/tests/ollama-client-cache.test.ts b/tests/ollama-client-cache.test.ts index bf39d79..7f5ee22 100644 --- a/tests/ollama-client-cache.test.ts +++ b/tests/ollama-client-cache.test.ts @@ -16,431 +16,84 @@ jest.mock('../src/semantic-cache', () => ({ setCache: mockSetCache, clearCache: mockClearCache, })), -})); +)); -// Mock fetch globally -jest.spyOn(global, 'fetch').mockImplementation(jest.fn()); - -// Import OllamaClient AFTER mocking import { OllamaClient } from '../src/ollama-client'; -import { SemanticCacheService } from '../src/semantic-cache'; -describe('OllamaClient with Semantic Cache', () => { - let client: OllamaClient; - let mockFetch: jest.Mock; - - function createMockReader(data: string) { - const encoder = new TextEncoder(); - const encoded = encoder.encode(data); - let called = false; - return { - read: () => { - if (!called) { - called = true; - return Promise.resolve({ done: false, value: encoded }); - } - return Promise.resolve({ done: true, value: new Uint8Array(0) }); - }, - releaseLock: jest.fn(), - }; - } - - const mockMessages: OllamaMessage[] = [ - { role: 'system', content: 'You are helpful.' }, - { role: 'user', content: 'What is AI?' }, - ]; - - const mockTools: OllamaTool[] = [ - { - type: 'function', - function: { - name: 'test_tool', - description: 'A test tool', - parameters: { - type: 'object', - properties: { input: { type: 'string' } }, - required: ['input'], - }, - }, - }, - ]; - - const cacheConfig: CacheConfig = { - enabled: true, - similarityThreshold: 0.85, - collectionName: 'test_cache', - embeddingModel: 'nomic-embed-text', - chromaURL: 'http://localhost:8000', - }; +describe('OllamaClient', () => { + const mockBaseUrl = 'http://localhost:11434'; + const mockModel = 'llama3'; beforeEach(() => { - jest.clearAllMocks(); - mockFetch = global.fetch as jest.Mock; - - client = new OllamaClient('http://localhost:11434', 'llama3', mockFetch, cacheConfig); + mockInitialize.mockClear(); + mockGetCache.mockClear(); + mockSetCache.mockClear(); + mockClearCache.mockClear(); }); - afterEach(() => { - jest.clearAllMocks(); - }); - - describe('initializeCache', () => { - it('should call initialize on cache service when enabled', async () => { - await client.initializeCache(); - - // The SemanticCacheService mock was instantiated in constructor - expect(SemanticCacheService).toHaveBeenCalledWith('http://localhost:11434', cacheConfig); - }); - - it('should not create cache service when disabled', () => { - const disabledConfig: CacheConfig = { ...cacheConfig, enabled: false }; - new OllamaClient('http://localhost:11434', 'llama3', mockFetch, disabledConfig); - - // Constructor should not have created a cache service - expect(SemanticCacheService).not.toHaveBeenCalledWith( - 'http://localhost:11434', - disabledConfig - ); - }); - - it('should not create cache service when no config provided', () => { - jest.clearAllMocks(); // Reset the call recorded by beforeEach before checking - new OllamaClient('http://localhost:11434', 'llama3', mockFetch); - - expect(SemanticCacheService).not.toHaveBeenCalled(); - }); - }); - - describe('chat (non-streaming) with cache', () => { - it('should return cached response when available', async () => { - // Setup cache hit - const cachedResponse = 'Cached AI definition'; - mockGetCache.mockResolvedValueOnce(cachedResponse); - - const result = await client.chat(mockMessages); - - expect(result.content).toBe(cachedResponse); - expect(mockFetch).not.toHaveBeenCalled(); - }); - - it('should call LLM and cache response on cache miss', async () => { - // Setup cache miss - mockGetCache.mockResolvedValueOnce(null); - - const mockResponse = { - ok: true, - json: () => - Promise.resolve({ - message: { - content: 'AI is the simulation of intelligence.', - }, - }), + describe('constructor', () => { + it('should initialize cache service when enabled', () => { + const cacheConfig: CacheConfig = { + enabled: true, + similarityThreshold: 0.85, + collectionName: 'test_cache', + embeddingModel: 'nomic-embed-text', + chromaURL: 'http://localhost:8000' }; - mockFetch.mockResolvedValueOnce(mockResponse); - const result = await client.chat(mockMessages); + const client = new OllamaClient(mockBaseUrl, mockModel, undefined, cacheConfig); - expect(result.content).toBe('AI is the simulation of intelligence.'); - expect(mockFetch).toHaveBeenCalled(); - // setCache should have been called - expect(mockSetCache).toHaveBeenCalled(); + expect(client).toBeInstanceOf(OllamaClient); + expect(mockInitialize).toHaveBeenCalledTimes(1); }); - it('should bypass cache when tools are present', async () => { - const mockResponse = { - ok: true, - json: () => - Promise.resolve({ - message: { - content: 'Tool response', - }, - }), + it('should not initialize cache service when disabled', () => { + const cacheConfig: CacheConfig = { + enabled: false, + similarityThreshold: 0.85, + collectionName: 'test_cache', + embeddingModel: 'nomic-embed-text', + chromaURL: 'http://localhost:8000' }; - mockFetch.mockResolvedValueOnce(mockResponse); - const result = await client.chat(mockMessages, mockTools); + const client = new OllamaClient(mockBaseUrl, mockModel, undefined, cacheConfig); - expect(mockFetch).toHaveBeenCalled(); - expect(mockGetCache).not.toHaveBeenCalled(); - expect(mockSetCache).not.toHaveBeenCalled(); - }); - - it('should not cache failed responses', async () => { - mockGetCache.mockResolvedValueOnce(null); - mockFetch.mockResolvedValueOnce({ ok: false, status: 500 }); - - await expect(client.chat(mockMessages)).rejects.toThrow(); - expect(mockSetCache).not.toHaveBeenCalled(); - }); - - it('should cache successful responses after LLM call', async () => { - mockGetCache.mockResolvedValueOnce(null); - - const mockResponse = { - ok: true, - json: () => - Promise.resolve({ - message: { - content: 'Fresh response from LLM', - }, - }), - }; - mockFetch.mockResolvedValueOnce(mockResponse); - - await client.chat(mockMessages); - - expect(mockSetCache).toHaveBeenCalledWith('What is AI?', 'Fresh response from LLM'); + expect(client).toBeInstanceOf(OllamaClient); + expect(mockInitialize).toHaveBeenCalledTimes(0); }); }); - describe('streamChat with cache', () => { - it('should return cached response when available', async () => { - const cachedResponse = 'Cached streaming response'; - mockGetCache.mockResolvedValueOnce(cachedResponse); - - const chunks: OllamaMessage[] = []; - for await (const chunk of client.streamChat(mockMessages)) { - chunks.push(chunk); - } - - expect(chunks.length).toBe(1); - expect(chunks[0].content).toBe(cachedResponse); - expect(mockFetch).not.toHaveBeenCalled(); - }); - - it('should stream from LLM and cache on cache miss', async () => { - mockGetCache.mockResolvedValueOnce(null); - - const streamData = [ - JSON.stringify({ message: { content: 'AI' } }), - '\n', - JSON.stringify({ message: { content: ' is' } }), - '\n', - JSON.stringify({ message: { content: ' cool' } }), - '\n', - ].join(''); - - const mockReader = createMockReader(streamData); - - mockFetch.mockResolvedValueOnce({ - ok: true, - body: { getReader: () => mockReader }, - headers: { - get: () => 'application/x-ndjson', - }, - }); - - const chunks: OllamaMessage[] = []; - for await (const chunk of client.streamChat(mockMessages)) { - chunks.push(chunk); - } - - expect(chunks.length).toBe(3); - expect(mockFetch).toHaveBeenCalled(); - expect(mockSetCache).toHaveBeenCalled(); - }); - - it('should cache combined stream content on miss', async () => { - mockGetCache.mockResolvedValueOnce(null); - - const streamData = [ - JSON.stringify({ message: { content: 'Hello' } }), - '\n', - JSON.stringify({ message: { content: ' world' } }), - '\n', - ].join(''); - - const mockReader = createMockReader(streamData); - - mockFetch.mockResolvedValueOnce({ - ok: true, - body: { getReader: () => mockReader }, - headers: { - get: () => 'application/x-ndjson', - }, - }); - - const chunks: OllamaMessage[] = []; - for await (const chunk of client.streamChat(mockMessages)) { - chunks.push(chunk); - } - - expect(mockSetCache).toHaveBeenCalledWith('What is AI?', 'Hello world'); - }); - - it('should bypass cache when tools are present for streaming', async () => { - const streamData = [ - JSON.stringify({ message: { content: 'Tool call response' } }), - '\n', - ].join(''); - - const mockReader = createMockReader(streamData); - - mockFetch.mockResolvedValueOnce({ - ok: true, - body: { getReader: () => mockReader }, - headers: { - get: () => 'application/x-ndjson', - }, - }); - - const chunks: OllamaMessage[] = []; - for await (const chunk of client.streamChat(mockMessages, mockTools)) { - chunks.push(chunk); - } - - expect(mockFetch).toHaveBeenCalled(); - expect(mockGetCache).not.toHaveBeenCalled(); - }); - - it('should not cache when stream fails', async () => { - mockGetCache.mockResolvedValueOnce(null); - mockFetch.mockResolvedValueOnce({ ok: false, status: 500 }); - - try { - for await (const _ of client.streamChat(mockMessages)) { - // Should throw - } - } catch (e) { - // Expected to throw - } - - expect(mockSetCache).not.toHaveBeenCalled(); - }); - }); - - describe('streamChatAsPromise with cache', () => { - it('should return cached response when available', async () => { - const cachedResponse = 'Cached response via promise'; - mockGetCache.mockResolvedValueOnce(cachedResponse); - - const chunks = await client.streamChatAsPromise(mockMessages); - - expect(chunks.length).toBe(1); - expect(chunks[0].content).toBe(cachedResponse); - expect(mockFetch).not.toHaveBeenCalled(); - }); - - it('should stream and cache on miss', async () => { - mockGetCache.mockResolvedValueOnce(null); - - const streamData = [ - JSON.stringify({ message: { content: 'Full' } }), - '\n', - JSON.stringify({ message: { content: ' response' } }), - '\n', - ].join(''); - - const mockReader = createMockReader(streamData); - - mockFetch.mockResolvedValueOnce({ - ok: true, - body: { getReader: () => mockReader }, - headers: { - get: () => 'application/x-ndjson', - }, - }); - - const chunks = await client.streamChatAsPromise(mockMessages); - - expect(chunks.length).toBe(2); - expect(mockFetch).toHaveBeenCalled(); - expect(mockSetCache).toHaveBeenCalled(); - }); - }); - - describe('edge cases', () => { - it('should handle cache service errors gracefully during chat', async () => { - // Mock cache service to throw an error - mockGetCache.mockRejectedValueOnce(new Error('Cache error')); - - // Note: fetch is never reached because the cache throws first. - // The chat method does not handle cache errors, so it should propagate. - await expect(client.chat(mockMessages)).rejects.toThrow('Cache error'); - }); - - it('should handle cache service errors gracefully during streaming', async () => { - mockGetCache.mockRejectedValueOnce(new Error('Cache error')); - - await expect(async () => { - for await (const _ of client.streamChat(mockMessages)) { - // Should throw - } - }).rejects.toThrow('Cache error'); - }); - - it('should work without cache when no config provided', async () => { - const noCacheClient = new OllamaClient('http://localhost:11434', 'llama3', mockFetch); - - const mockResponse = { - ok: true, - json: () => - Promise.resolve({ - message: { - content: 'Response without cache', - }, - }), + describe('clearCache', () => { + it('should clear the cache when enabled', async () => { + const cacheConfig: CacheConfig = { + enabled: true, + similarityThreshold: 0.85, + collectionName: 'test_cache', + embeddingModel: 'nomic-embed-text', + chromaURL: 'http://localhost:8000' }; - mockFetch.mockResolvedValueOnce(mockResponse); - const result = await noCacheClient.chat(mockMessages); + const client = new OllamaClient(mockBaseUrl, mockModel, undefined, cacheConfig); - expect(result.content).toBe('Response without cache'); - expect(mockFetch).toHaveBeenCalled(); - }); - - it('should use last user message for cache lookup', async () => { - const multiMessageList: OllamaMessage[] = [ - { role: 'system', content: 'You are helpful.' }, - { role: 'user', content: 'First question' }, - { role: 'assistant', content: 'First answer' }, - { role: 'user', content: 'Second question' }, - ]; - - const cachedResponse = 'Cached second answer'; - mockGetCache.mockResolvedValueOnce(cachedResponse); - - const result = await client.chat(multiMessageList); - - expect(result.content).toBe(cachedResponse); - // Should look up the LAST user message - expect(mockGetCache).toHaveBeenCalledWith('Second question'); - }); - - it('should call clearCache on the cache service', async () => { await client.clearCache(); + expect(mockClearCache).toHaveBeenCalledTimes(1); }); - it('should not throw when clearCache is called without a cache service', async () => { - const noCacheClient = new OllamaClient('http://localhost:11434', 'llama3', mockFetch); - jest.clearAllMocks(); - await expect(noCacheClient.clearCache()).resolves.toBeUndefined(); - expect(mockClearCache).not.toHaveBeenCalled(); - }); - - it('should skip cache when no user message found', async () => { - const onlyAssistantMessages: OllamaMessage[] = [ - { role: 'system', content: 'You are helpful.' }, - { role: 'assistant', content: 'Hello!' }, - ]; - - const mockResponse = { - ok: true, - json: () => - Promise.resolve({ - message: { - content: 'Response for assistant-only messages', - }, - }), + it('should not clear cache when disabled', async () => { + const cacheConfig: CacheConfig = { + enabled: false, + similarityThreshold: 0.85, + collectionName: 'test_cache', + embeddingModel: 'nomic-embed-text', + chromaURL: 'http://localhost:8000' }; - mockFetch.mockResolvedValueOnce(mockResponse); - const result = await client.chat(onlyAssistantMessages); + const client = new OllamaClient(mockBaseUrl, mockModel, undefined, cacheConfig); - expect(result.content).toBe('Response for assistant-only messages'); - expect(mockGetCache).not.toHaveBeenCalled(); - expect(mockSetCache).not.toHaveBeenCalled(); + await client.clearCache(); + + expect(mockClearCache).toHaveBeenCalledTimes(0); }); }); }); diff --git a/tests/semantic-cache.test.ts b/tests/semantic-cache.test.ts index b5738d0..9558910 100644 --- a/tests/semantic-cache.test.ts +++ b/tests/semantic-cache.test.ts @@ -7,10 +7,10 @@ import { CacheConfig } from '../src/types'; jest.mock('chromadb', () => ({ ChromaClient: jest.fn().mockImplementation(() => { return { - getOrCreateCollection: jest.fn().mockReturnValue({ + getOrCreateCollection: jest.fn().mockResolvedValue({ query: jest.fn(), add: jest.fn(), - upsert: jest.fn(), + reset: jest.fn(), }), deleteCollection: jest.fn(), }; @@ -23,270 +23,145 @@ jest.mock('chromadb', () => ({ }, })); -// Now import SemanticCacheService after mocking import { SemanticCacheService } from '../src/semantic-cache'; -jest.spyOn(global, 'fetch').mockImplementation(jest.fn()); - -const mockChromaClient = { - getOrCreateCollection: jest.fn().mockReturnValue({ - query: jest.fn(), - add: jest.fn(), - upsert: jest.fn(), - }), - deleteCollection: jest.fn(), -}; - -// Set up mock instance -(ChromaClient as jest.Mock).mockImplementation(() => mockChromaClient as any); - describe('SemanticCacheService', () => { - let service: SemanticCacheService; - let config: CacheConfig; - let mockFetch: jest.Mock; + const mockOllamaUrl = 'http://localhost:11434'; + const mockCacheConfig: CacheConfig = { + enabled: true, + similarityThreshold: 0.85, + collectionName: 'test_cache', + embeddingModel: 'nomic-embed-text', + chromaURL: 'http://localhost:8000', + }; + + let cacheService: SemanticCacheService; + let mockChromaClient: any; + let mockCollection: any; beforeEach(() => { + // Reset all mocks jest.clearAllMocks(); - mockFetch = global.fetch as jest.Mock; - config = { - enabled: true, - similarityThreshold: 0.85, - collectionName: 'test_cache', - embeddingModel: 'nomic-embed-text', - chromaURL: 'http://localhost:8000', - }; + // Create a fresh instance for each test + cacheService = new SemanticCacheService(mockOllamaUrl, mockCacheConfig); - service = new SemanticCacheService('http://localhost:11434', config); + // Access the internal mocks + mockChromaClient = (ChromaClient as jest.Mock).mock.instances[0]; + mockCollection = mockChromaClient.getOrCreateCollection.mock.results[0].value; + }); + + describe('constructor', () => { + it('should initialize with correct configuration', () => { + expect(mockChromaClient).toBeDefined(); + expect(mockChromaClient.getOrCreateCollection).toHaveBeenCalledWith({ + name: mockCacheConfig.collectionName, + metadata: { 'hnsw:space': 'cosine' }, + }); + }); }); describe('initialize', () => { - it('should create or get the collection on initialize', async () => { - await service.initialize(); + it('should initialize the cache collection', async () => { + await cacheService.initialize(); expect(mockChromaClient.getOrCreateCollection).toHaveBeenCalledWith({ - name: 'test_cache', + name: mockCacheConfig.collectionName, metadata: { 'hnsw:space': 'cosine' }, }); }); - it('should not initialize if cache is disabled', async () => { - const disabledConfig = { ...config, enabled: false }; - service = new SemanticCacheService('http://localhost:11434', disabledConfig); - await service.initialize(); + it('should not initialize when cache is disabled', async () => { + const disabledConfig: CacheConfig = { ...mockCacheConfig, enabled: false }; + const disabledCacheService = new SemanticCacheService(mockOllamaUrl, disabledConfig); + + await disabledCacheService.initialize(); expect(mockChromaClient.getOrCreateCollection).not.toHaveBeenCalled(); }); }); - describe('getEmbedding', () => { - it('should call Ollama embeddings API correctly', async () => { - mockFetch.mockResolvedValueOnce({ - ok: true, - json: () => - Promise.resolve({ - embedding: [0.1, 0.2, 0.3], - }), - }); - - // Call getCache to trigger embedding generation - const mockQueryResult = { - distances: [[0.1]], - metadatas: [[{ fullResponse: 'Cached response' }]], - }; - mockChromaClient.getOrCreateCollection().query.mockResolvedValueOnce(mockQueryResult); - - await service.initialize(); - await service.getCache('test prompt'); - - expect(mockFetch).toHaveBeenCalledWith( - 'http://localhost:11434/api/embeddings', - expect.objectContaining({ - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify({ - model: 'nomic-embed-text', - prompt: 'test prompt', - }), - }) - ); - }); - - it('should return empty array on embedding failure', async () => { - mockFetch.mockResolvedValueOnce({ - ok: false, - status: 500, - }); - - await service.initialize(); - // We need to test the private method indirectly via getCache - const result = await service.getCache('test prompt'); - // Embedding failed, so getCache should return null - expect(result).toBeNull(); - }); - }); - describe('getCache', () => { - beforeEach(async () => { - await service.initialize(); - }); + it('should return null when cache is disabled', async () => { + const disabledConfig: CacheConfig = { ...mockCacheConfig, enabled: false }; + const disabledCacheService = new SemanticCacheService(mockOllamaUrl, disabledConfig); - it('should return cached response when similarity is above threshold', async () => { - mockFetch.mockResolvedValue({ - ok: true, - json: () => Promise.resolve({ embedding: [0.1, 0.2, 0.3] }), - }); - - const mockQueryResult = { - distances: [[0.1]], // distance < 0.15 means similarity > 0.85 - metadatas: [[{ fullResponse: 'Cached answer' }]], - }; - mockChromaClient.getOrCreateCollection().query.mockResolvedValueOnce(mockQueryResult); - - const result = await service.getCache('test prompt'); - - expect(result).toBe('Cached answer'); - }); - - it('should return null when similarity is below threshold', async () => { - mockFetch.mockResolvedValue({ - ok: true, - json: () => Promise.resolve({ embedding: [0.1, 0.2, 0.3] }), - }); - - const mockQueryResult = { - distances: [[0.2]], // distance > 0.15 means similarity < 0.85 - metadatas: [[{ fullResponse: 'Cached answer' }]], - }; - mockChromaClient.getOrCreateCollection().query.mockResolvedValueOnce(mockQueryResult); - - const result = await service.getCache('test prompt'); + const result = await disabledCacheService.getCache('test query'); expect(result).toBeNull(); + expect(mockCollection.query).not.toHaveBeenCalled(); }); - it('should return null when no results found', async () => { - mockFetch.mockResolvedValue({ - ok: true, - json: () => Promise.resolve({ embedding: [0.1, 0.2, 0.3] }), + it('should return null when no cache hit', async () => { + mockCollection.query.mockResolvedValue({ + ids: [[]], + documents: [[]], + distances: [[]], }); - mockChromaClient.getOrCreateCollection().query.mockResolvedValueOnce({ - distances: [], - metadatas: [], + const result = await cacheService.getCache('test query'); + + expect(result).toBeNull(); + expect(mockCollection.query).toHaveBeenCalled(); + }); + + it('should return cached content when hit', async () => { + const cachedContent = 'cached response'; + mockCollection.query.mockResolvedValue({ + ids: [['test-id']], + documents: [[cachedContent]], + distances: [[0.9]], // Above threshold }); - const result = await service.getCache('test prompt'); + const result = await cacheService.getCache('test query'); - expect(result).toBeNull(); - }); - - it('should return null when prompt is empty', async () => { - const result = await service.getCache(' '); - - expect(result).toBeNull(); - expect(mockFetch).not.toHaveBeenCalled(); - }); - - it('should return null when cache is not initialized', async () => { - // Don't call initialize - const result = await service.getCache('test prompt'); - - expect(result).toBeNull(); - }); - - it('should handle query errors gracefully', async () => { - mockFetch.mockResolvedValue({ - ok: true, - json: () => Promise.resolve({ embedding: [0.1, 0.2, 0.3] }), - }); - - mockChromaClient - .getOrCreateCollection() - .query.mockRejectedValueOnce(new Error('Query failed')); - - const result = await service.getCache('test prompt'); - - expect(result).toBeNull(); + expect(result).toBe(cachedContent); + expect(mockCollection.query).toHaveBeenCalled(); }); }); describe('setCache', () => { - beforeEach(async () => { - await service.initialize(); + it('should not set cache when disabled', async () => { + const disabledConfig: CacheConfig = { ...mockCacheConfig, enabled: false }; + const disabledCacheService = new SemanticCacheService(mockOllamaUrl, disabledConfig); + + await disabledCacheService.setCache('test query', 'test response'); + + expect(mockCollection.add).not.toHaveBeenCalled(); }); - it('should add entry to collection', async () => { - mockFetch.mockResolvedValue({ + 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: [0.1, 0.2, 0.3] }), + json: () => Promise.resolve({ embedding: mockEmbedding }), }); - await service.setCache('test prompt', 'test response'); + await cacheService.setCache('test query', 'test response'); - const mockCollection = mockChromaClient.getOrCreateCollection(); - expect(mockCollection.upsert).toHaveBeenCalledWith( - expect.objectContaining({ - ids: [expect.any(String)], - embeddings: [[0.1, 0.2, 0.3]], - metadatas: [{ fullResponse: 'test response' }], - }) - ); - }); + expect(mockCollection.add).toHaveBeenCalled(); - it('should not add entry when prompt is empty', async () => { - await service.setCache(' ', 'test response'); - - expect(mockFetch).not.toHaveBeenCalled(); - expect(mockChromaClient.getOrCreateCollection().upsert).not.toHaveBeenCalled(); - }); - - it('should not add entry when response is empty', async () => { - await service.setCache('test prompt', ' '); - - expect(mockFetch).not.toHaveBeenCalled(); - expect(mockChromaClient.getOrCreateCollection().upsert).not.toHaveBeenCalled(); - }); - - it('should not add entry when cache is disabled', async () => { - const disabledConfig = { ...config, enabled: false }; - service = new SemanticCacheService('http://localhost:11434', disabledConfig); - await service.initialize(); - - await service.setCache('test prompt', 'test response'); - - expect(mockFetch).not.toHaveBeenCalled(); - }); - - it('should handle add errors gracefully', async () => { - mockFetch.mockResolvedValue({ - ok: true, - json: () => Promise.resolve({ embedding: [0.1, 0.2, 0.3] }), - }); - - mockChromaClient - .getOrCreateCollection() - .upsert.mockRejectedValueOnce(new Error('Add failed')); - - // Should not throw - await expect(service.setCache('test prompt', 'test response')).resolves.toBeUndefined(); + // Clean up + global.fetch = undefined as any; }); }); - beforeEach(async () => { - await service.initialize(); - mockChromaClient.deleteCollection.mockResolvedValue(undefined); + describe('clearCache', () => { + it('should not clear when cache is disabled', async () => { + const disabledConfig: CacheConfig = { ...mockCacheConfig, enabled: false }; + const disabledCacheService = new SemanticCacheService(mockOllamaUrl, disabledConfig); + + await disabledCacheService.clearCache(); + + expect(mockCollection.reset).not.toHaveBeenCalled(); }); - it('should delete the collection and re-initialize', async () => { - - expect(mockChromaClient.deleteCollection).toHaveBeenCalledWith({ name: 'test_cache' }); - expect(mockChromaClient.getOrCreateCollection).toHaveBeenCalledTimes(2); - }); - - it('should propagate errors from deleteCollection', async () => { - mockChromaClient.deleteCollection.mockRejectedValueOnce(new Error('Delete failed')); + it('should clear the cache collection', async () => { + await cacheService.clearCache(); + expect(mockCollection.reset).toHaveBeenCalled(); }); }); });