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
This commit is contained in:
+34
-25
@@ -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
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
|
||||
+8
-77
@@ -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<string, string> = 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<void> {
|
||||
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<string> {
|
||||
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;
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user