From 0a0173363bca29ad052b25b0e525ac95c30a0ec3 Mon Sep 17 00:00:00 2001 From: Florian Egger Date: Wed, 6 May 2026 19:58:26 +0200 Subject: [PATCH] Add MAX_STREAM_CHUNKS constant for streaming control Improve path validation to detect parent directory traversal Refactor VaultIndexer to use VaultLike interface Enhance heading matching to support 1-6 level headings Improve token stemming with length-based checks Optimize context building for user messages --- src/chat-view.ts | 59 +++++++++++++++++------------- src/tool-executor.ts | 2 +- src/vault-indexer.ts | 85 +++++--------------------------------------- 3 files changed, 43 insertions(+), 103 deletions(-) diff --git a/src/chat-view.ts b/src/chat-view.ts index d478df6..2618896 100755 --- a/src/chat-view.ts +++ b/src/chat-view.ts @@ -8,6 +8,7 @@ type MouseEvent = globalThis.MouseEvent; const DEFAULT_VAULT_SEARCH_LIMIT = 3; const MAX_MESSAGE_HISTORY = 50; +const MAX_STREAM_CHUNKS = 1000; import { PluginSettings, OllamaMessage, @@ -204,7 +205,10 @@ export class ChatView extends ItemView { void this.inputKeyDownHandler?.(e); }; (this.sendButton as HTMLButtonElement).addEventListener('click', this.sendButtonEventHandler); - (this.inputEl as HTMLTextAreaElement).addEventListener('keydown', this.inputKeyDownEventHandler); + (this.inputEl as HTMLTextAreaElement).addEventListener( + 'keydown', + this.inputKeyDownEventHandler + ); if (this.newChatButton) { if (!this.newChatButtonClickHandler) { this.newChatButtonClickHandler = () => this.clearConversation(); @@ -288,9 +292,7 @@ export class ChatView extends ItemView { // Search vault using user message as query const entries = await this.vaultIndexer.searchVault(userMessage, DEFAULT_VAULT_SEARCH_LIMIT); - let context = entries - .map((entry) => `### ${entry.title}\n${entry.content}`) - .join('\n\n'); + 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; @@ -302,9 +304,10 @@ export class ChatView extends ItemView { role: 'system', content: 'You are a helpful assistant.', }; + const userContent = context ? `${context}\n\n${userMessage}` : userMessage; const userMessageWithContext: OllamaMessage = { role: 'user', - content: `${context}\n\n${userMessage}`, + content: userContent, }; const messages: OllamaMessage[] = [ @@ -329,8 +332,11 @@ export class ChatView extends ItemView { parameters: { type: 'object' as const, properties: { - path: { type: 'string' as const }, - content: { type: 'string' as const }, + path: { + type: 'string' as const, + description: "Relative path within the vault, e.g. 'Notes/todo.md'", + }, + content: { type: 'string' as const, description: 'Content of the file to create' }, }, required: ['path', 'content'], }, @@ -369,12 +375,9 @@ export class ChatView extends ItemView { let fullResponse = ''; let toolCalls: ToolCall[] = []; let chunkCount = 0; - const MAX_STREAM_CHUNKS = 1000; - const maxChunks = MAX_STREAM_CHUNKS; - for await (const chunk of stream) { chunkCount++; - if (chunkCount > maxChunks) { + if (chunkCount > MAX_STREAM_CHUNKS) { throw new Error('Response too long, stopped streaming'); } @@ -418,22 +421,28 @@ export class ChatView extends ItemView { } } - // 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), - })), - ]; + // 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); + 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 }); + // Update the assistant message with the final response immutably + this.updateMessageById(assistantMessageId, { + content: fullResponse, + isStreaming: false, + }); + } } // Update last message immutably — only if no tool calls were processed diff --git a/src/tool-executor.ts b/src/tool-executor.ts index de6ac64..c81bb11 100644 --- a/src/tool-executor.ts +++ b/src/tool-executor.ts @@ -51,7 +51,7 @@ export class ToolExecutor { // Reject paths that traverse to parent directories const normalized = path.replace(/^(\.\/)+/, ''); - if (normalized.includes('../')) { + if (normalized.split('/').includes('..')) { return false; } diff --git a/src/vault-indexer.ts b/src/vault-indexer.ts index 0c2aaa4..be5f0e6 100644 --- a/src/vault-indexer.ts +++ b/src/vault-indexer.ts @@ -1,8 +1,6 @@ // src/vault-indexer.ts -import { OllamaClient } from './ollama-client'; -import { ValidationError, VaultIndexEntry } from './types'; -import { ErrorHandler } from './error-handler'; +import { VaultIndexEntry } from './types'; import { Logger, sanitizeFilePath } from './utils'; interface Frontmatter { @@ -44,80 +42,11 @@ function isVaultLike(value: unknown): value is VaultLike { } class VaultIndexer { - private ollamaClient: OllamaClient | null = null; private summaries: Map = new Map(); private vault: VaultLike | null = null; - constructor(vaultOrClient: VaultLike | OllamaClient | null) { - // Support both old (OllamaClient) and new (VaultLike) interfaces - if (isVaultLike(vaultOrClient)) { - this.vault = vaultOrClient; - } else { - this.ollamaClient = vaultOrClient; - } - } - - async indexVault(vaultPath: string): Promise { - try { - const files = this.getMarkdownFilesInVault(vaultPath); - for (const file of files) { - const content = this.readFileContent(file.path); - const summary = await this.summarizeFile(content); - this.storeSummary(file.path, summary); - } - } catch (error: unknown) { - if (error instanceof Error) { - ErrorHandler.handleError(error, 'VaultIndexer.indexVault'); - } else { - throw new ValidationError('An unexpected error occurred while indexing the vault'); - } - } - } - - private getMarkdownFilesInVault(vaultPath: string): { path: string }[] { - // Simulate getting markdown files from the vault - const sanitizedPath = sanitizeFilePath(vaultPath); - try { - // This is a placeholder for actual file system operations - return [{ path: `${sanitizedPath}/file1.md` }, { path: `${sanitizedPath}/file2.md` }]; - } catch (error) { - throw new ValidationError('Failed to get markdown files from vault'); - } - } - - private readFileContent(filePath: string): string { - const sanitizedPath = sanitizeFilePath(filePath); - try { - // This is a placeholder for actual file reading operations - return `Content of ${sanitizedPath}`; - } catch (error) { - throw new ValidationError('Failed to read file content'); - } - } - - private async summarizeFile(content: string): Promise { - if (!this.ollamaClient) { - throw new ValidationError('OllamaClient not available for summarization'); - } - - try { - // Use the existing chat API to summarize text - const messages = [ - { - role: 'system' as const, - content: 'Summarize the following text concisely:', - }, - { - role: 'user' as const, - content: content, - }, - ]; - - const response = await this.ollamaClient.chat(messages); - return response.content; - } catch (error) { - throw new ValidationError('Failed to summarize file'); - } + constructor(vault: VaultLike) { + this.vault = vault; } private storeSummary(filePath: string, summary: string): void { @@ -268,7 +197,7 @@ class VaultIndexer { } } - const headingMatches = content.match(/^# (.*?)$/gm); + const headingMatches = content.match(/^#{1,6} (.*?)$/gm); if (headingMatches) { headings.push(...headingMatches.map((h: string) => h.replace(/^# /, ''))); } @@ -350,9 +279,11 @@ class VaultIndexer { } private stemToken(token: string): string { + // Improved stemmer that handles edge cases + 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')) return token.slice(0, -2); - if (token.endsWith('ing')) return token.slice(0, -3); + 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; }