Merge branch 'modularize-indexing-pipeline'
This commit is contained in:
+433
-404
@@ -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;
|
||||
|
||||
+354
-317
@@ -1,58 +1,41 @@
|
||||
import { ItemView, WorkspaceLeaf, Notice } from 'obsidian';
|
||||
/// <reference lib="dom" />
|
||||
// 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<void>) | null = null;
|
||||
private inputKeyDownHandler: ((e: KeyboardEvent) => Promise<void>) | 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<void>) | null {
|
||||
getSendButtonClickHandler() {
|
||||
return this.sendButtonClickHandler;
|
||||
}
|
||||
|
||||
public getInputKeyDownHandler(): ((e: KeyboardEvent) => Promise<void>) | 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<void> {
|
||||
async clearCache(): Promise<void> {
|
||||
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<void> {
|
||||
async onClose(): Promise<void> {
|
||||
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<ChatMessage>): boolean {
|
||||
updateMessageById(id: string, updates: Partial<ChatMessage>): 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<ChatMessage>): 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<void> {
|
||||
// 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<void> {
|
||||
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;
|
||||
|
||||
+14
-14
@@ -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',
|
||||
},
|
||||
};
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
// Default plugin settings
|
||||
|
||||
export const DEFAULT_SETTINGS = {
|
||||
ollamaUrl: 'http://localhost:11434',
|
||||
model: 'llama3',
|
||||
|
||||
@@ -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;
|
||||
@@ -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();
|
||||
}
|
||||
}
|
||||
@@ -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; } });
|
||||
@@ -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';
|
||||
@@ -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;
|
||||
@@ -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<string, unknown>;
|
||||
firstParagraph?: string;
|
||||
}
|
||||
|
||||
interface NormalizedContent {
|
||||
path: string;
|
||||
title: string;
|
||||
content: string;
|
||||
tokens: string[];
|
||||
headings: string[];
|
||||
frontmatter: Record<string, unknown>;
|
||||
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<string, unknown>): Record<string, unknown> {
|
||||
const normalized: Record<string, unknown> = {};
|
||||
|
||||
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<string, unknown>, 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<string, unknown>;
|
||||
firstParagraph?: string;
|
||||
wordCount: number;
|
||||
createdAt?: string;
|
||||
updatedAt?: string;
|
||||
// Additional enrichment fields for vectorization
|
||||
chunkIndex: number;
|
||||
chunkSize: number;
|
||||
}
|
||||
@@ -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;
|
||||
@@ -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<VaultIndexEntry | null> {
|
||||
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<string, string>,
|
||||
batchSize: number = 10
|
||||
): Promise<VaultIndexEntry[]> {
|
||||
const results: VaultIndexEntry[] = [];
|
||||
const seenPaths = new Set<string>();
|
||||
|
||||
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<typeof result> => result !== null
|
||||
);
|
||||
results.push(...validResults);
|
||||
}
|
||||
|
||||
return results;
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
@@ -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<number[]> {
|
||||
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');
|
||||
}
|
||||
}
|
||||
+191
-186
@@ -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();
|
||||
}
|
||||
}
|
||||
|
||||
+108
-114
@@ -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<PluginSettings> | 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<void> {
|
||||
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();
|
||||
}
|
||||
|
||||
+163
-283
@@ -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;
|
||||
|
||||
+71
-229
@@ -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<OllamaMessage>;
|
||||
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<OllamaMessage[]> {
|
||||
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<OllamaMessage> {
|
||||
// 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<OllamaMessage, void, unknown> {
|
||||
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<void>((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<string, unknown>;
|
||||
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<string, unknown>;
|
||||
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<OllamaMessage> {
|
||||
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<string, unknown>): 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<OllamaMessage>;
|
||||
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,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
+85
-107
@@ -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;
|
||||
|
||||
+58
-80
@@ -6,20 +6,19 @@ import { CacheConfig } from './types';
|
||||
|
||||
export class SemanticCacheService {
|
||||
private client: ChromaClient;
|
||||
private collection: any | null = null;
|
||||
private collection: ReturnType<ChromaClient['getOrCreateCollection']> | 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<void> {
|
||||
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<number[]> {
|
||||
async getCache(query: string): Promise<string | null> {
|
||||
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<string | null> {
|
||||
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<void> {
|
||||
if (!this.collection || !this.config.enabled || !prompt.trim() || !response.trim()) {
|
||||
return;
|
||||
}
|
||||
async setCache(query: string, response: string): Promise<void> {
|
||||
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<void> {
|
||||
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<number[]> {
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
+38
-106
@@ -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<string, unknown>;
|
||||
required: string[];
|
||||
};
|
||||
};
|
||||
}
|
||||
|
||||
export interface OllamaToolCall {
|
||||
id: string;
|
||||
type: 'function';
|
||||
function: {
|
||||
name: string;
|
||||
arguments: string | Record<string, unknown>;
|
||||
};
|
||||
}
|
||||
|
||||
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<string, unknown>;
|
||||
};
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
|
||||
+140
-257
@@ -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());
|
||||
}
|
||||
|
||||
+117
-265
@@ -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<string | null>;
|
||||
put(key: string, value: string): Promise<void>;
|
||||
clear(): Promise<void>;
|
||||
}
|
||||
|
||||
class InMemoryCache implements Cache {
|
||||
private store: Map<string, string>;
|
||||
|
||||
constructor() {
|
||||
this.store = new Map();
|
||||
}
|
||||
|
||||
get(key: string): Promise<string | null> {
|
||||
return Promise.resolve(this.store.get(key) || null);
|
||||
}
|
||||
|
||||
put(key: string, value: string): Promise<void> {
|
||||
this.store.set(key, value);
|
||||
return Promise.resolve();
|
||||
}
|
||||
|
||||
clear(): Promise<void> {
|
||||
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<string>;
|
||||
}
|
||||
|
||||
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<VaultIndexEntry[]> {
|
||||
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<VaultIndexEntry[]> {
|
||||
const batchSize = 10;
|
||||
const results: VaultIndexEntry[] = [];
|
||||
const seenPaths = new Set<string>();
|
||||
|
||||
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<typeof result> => 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<string> = new Set<string>();
|
||||
|
||||
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());
|
||||
}
|
||||
|
||||
@@ -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');
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
+91
-216
@@ -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();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user