Refactor Ollama client and related components
Update error handling and improve streaming capabilities in the Ollama client and related components. Key changes include: - Simplify error types and improve error handling - Refactor streaming logic to use async generators - Update tool execution and vault indexing - Improve utility functions and types - Update test files to reflect changes
This commit is contained in:
Regular → Executable
Regular → Executable
Regular → Executable
Regular → Executable
Regular → Executable
+29
-12
@@ -1,32 +1,49 @@
|
||||
// Mock for ollama-client for testing
|
||||
import { OllamaMessage, OllamaTool, ToolCall } from '../src/types';
|
||||
import { OllamaMessage, ToolCall } from '../src/types';
|
||||
import type { APIError } from '../src/error-handler';
|
||||
|
||||
export class OllamaClient {
|
||||
private url: string;
|
||||
private model: string;
|
||||
|
||||
// Mock fetch function for testing
|
||||
private fetchFn: typeof fetch = jest.fn();
|
||||
private fetchFn: jest.Mock<Promise<Response>, [string, RequestInit?]>> = jest.fn();
|
||||
|
||||
constructor(url: string, model: string, fetchFn?: typeof fetch) {
|
||||
this.url = url;
|
||||
this.model = model;
|
||||
if (fetchFn) this.fetchFn = fetchFn;
|
||||
if (fetchFn) this.fetchFn = jest.fn(fetchFn);
|
||||
}
|
||||
|
||||
async streamChatMessages(
|
||||
async *streamChatMessages(
|
||||
prompt: string,
|
||||
options: { abortSignal?: AbortSignal } = {}
|
||||
): Promise<string> {
|
||||
// Mock implementation - return a simple response
|
||||
return `Mock response for: ${prompt}`;
|
||||
): AsyncGenerator<OllamaMessage, void, unknown> {
|
||||
// Mock implementation - simulate streaming response
|
||||
const mockResponse = [
|
||||
{ role: 'assistant', content: 'Part 1' },
|
||||
{ role: 'assistant', content: 'Part 2' }
|
||||
];
|
||||
|
||||
for (const message of mockResponse) {
|
||||
yield message;
|
||||
await new Promise(resolve => setTimeout(resolve, 100)); // Simulate delay
|
||||
}
|
||||
}
|
||||
|
||||
async streamToolMessages(
|
||||
toolCall: string,
|
||||
async *streamToolMessages(
|
||||
toolCall: ToolCall,
|
||||
options: { abortSignal?: AbortSignal } = {}
|
||||
): Promise<string> {
|
||||
// Mock implementation - return a simple tool response
|
||||
return `Mock tool response for: ${toolCall}`;
|
||||
): AsyncGenerator<OllamaMessage, void, unknown> {
|
||||
// Mock implementation - simulate streaming response for tool
|
||||
const mockResponse = [
|
||||
{ role: 'assistant', content: `Tool ${toolCall.tool_name} Part 1` },
|
||||
{ role: 'assistant', content: `Tool ${toolCall.tool_name} Part 2` }
|
||||
];
|
||||
|
||||
for (const message of mockResponse) {
|
||||
yield message;
|
||||
await new Promise(resolve => setTimeout(resolve, 100)); // Simulate delay
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Regular → Executable
Regular → Executable
Regular → Executable
Regular → Executable
Regular → Executable
|
Before Width: | Height: | Size: 445 B After Width: | Height: | Size: 445 B |
Regular → Executable
Regular → Executable
Regular → Executable
Regular → Executable
Regular → Executable
|
Before Width: | Height: | Size: 138 B After Width: | Height: | Size: 138 B |
Regular → Executable
Regular → Executable
Regular → Executable
Regular → Executable
Regular → Executable
Regular → Executable
Regular → Executable
Regular → Executable
Regular → Executable
Regular → Executable
+15
-172
@@ -1,178 +1,21 @@
|
||||
import { Notice } from 'obsidian';
|
||||
import {
|
||||
OllamaError,
|
||||
ErrorType,
|
||||
NetworkError,
|
||||
ApiError,
|
||||
ValidationError,
|
||||
StreamingError,
|
||||
ToolExecutionError,
|
||||
PathValidationError,
|
||||
} from './types';
|
||||
// src/error-handler.ts
|
||||
|
||||
import { NetworkError, ApiError, UserInputError } from './errors';
|
||||
|
||||
export class ErrorHandler {
|
||||
/**
|
||||
* Centralized error handling for the Ollama plugin
|
||||
* Provides consistent error messages and logging
|
||||
*/
|
||||
static handleError(error: unknown, context?: string): void {
|
||||
let userMessage = 'An unexpected error occurred';
|
||||
let shouldShowError = true;
|
||||
|
||||
if (error instanceof OllamaError) {
|
||||
userMessage = this.getUserFriendlyMessage(error);
|
||||
shouldShowError = true;
|
||||
} else if (error instanceof Error) {
|
||||
userMessage = this.getUserFriendlyMessageFromError(error);
|
||||
shouldShowError = true;
|
||||
static handle(error: unknown): void {
|
||||
if (error instanceof NetworkError) {
|
||||
console.error('Network Error:', error.message);
|
||||
// Handle network errors, e.g., show a notification to the user
|
||||
} else if (error instanceof ApiError) {
|
||||
console.error('API Error:', error.message, 'Status Code:', error.statusCode);
|
||||
// Handle API errors, e.g., show a notification with status code
|
||||
} else if (error instanceof UserInputError) {
|
||||
console.warn('User Input Error:', error.message);
|
||||
// Handle user input errors, e.g., highlight the input field
|
||||
} else {
|
||||
userMessage = 'An unexpected error occurred';
|
||||
shouldShowError = true;
|
||||
console.error('Unexpected Error:', error);
|
||||
// Handle unexpected errors, e.g., log to a service or show a generic message
|
||||
}
|
||||
|
||||
if (shouldShowError) {
|
||||
new Notice(userMessage);
|
||||
}
|
||||
|
||||
// Log detailed error for debugging
|
||||
console.error(
|
||||
`[OllamaPlugin${context ? ' ' + context : ''}] ${error instanceof Error ? error.message : 'Unknown error'}`
|
||||
);
|
||||
if (error instanceof Error) {
|
||||
console.error('[Stack]', error.stack);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get user-friendly message from specific error types
|
||||
*/
|
||||
private static getUserFriendlyMessage(error: OllamaError): string {
|
||||
switch (error.type) {
|
||||
case ErrorType.NETWORK_ERROR:
|
||||
if (error instanceof NetworkError) {
|
||||
return 'Connection error. Please check if Ollama is running.';
|
||||
}
|
||||
return 'Network error. Please check your connection to Ollama.';
|
||||
|
||||
case ErrorType.API_ERROR:
|
||||
if (error instanceof ApiError) {
|
||||
return 'Ollama API error. Please check the Ollama logs for details.';
|
||||
}
|
||||
return 'API communication error. Please try again.';
|
||||
|
||||
case ErrorType.VALIDATION_ERROR:
|
||||
if (error instanceof ValidationError) {
|
||||
const details = error.validationDetails;
|
||||
if (details?.field) {
|
||||
return `Invalid ${details.field}. ${details.message || 'Please check your input.'}`;
|
||||
}
|
||||
return 'Input validation error. Please correct your input.';
|
||||
}
|
||||
return 'Input validation error. Please correct your input.';
|
||||
|
||||
case ErrorType.STREAMING_ERROR:
|
||||
if (error instanceof StreamingError) {
|
||||
return 'Response too long. Please try a shorter request.';
|
||||
}
|
||||
return 'Streaming error. Please try again.';
|
||||
|
||||
case ErrorType.TOOL_EXECUTION_ERROR:
|
||||
if (error instanceof ToolExecutionError) {
|
||||
return `Tool error: ${error.toolName || 'tool'} failed to execute. Please try again.`;
|
||||
}
|
||||
return 'Tool execution error. Please try a different command.';
|
||||
|
||||
case ErrorType.PATH_VALIDATION_ERROR:
|
||||
if (error instanceof PathValidationError) {
|
||||
return 'Invalid file path. Please use a relative path without special characters.';
|
||||
}
|
||||
return 'Path validation error. Please check your file path.';
|
||||
|
||||
case ErrorType.UNKNOWN_ERROR:
|
||||
return 'An unexpected error occurred. Please try again.';
|
||||
|
||||
default:
|
||||
return error.message || 'An error occurred';
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get user-friendly message from generic Error
|
||||
*/
|
||||
/**
|
||||
* Get user-friendly message from generic Error
|
||||
* Note: This method uses substring matching which is inherently fragile.
|
||||
* If an error message happens to contain certain keywords but isn't actually
|
||||
* that type of error, it may be misclassified. This heuristic approach
|
||||
* provides a good balance between robustness and accuracy for most common cases.
|
||||
*/
|
||||
private static getUserFriendlyMessageFromError(error: Error): string {
|
||||
const message = error.message.toLowerCase();
|
||||
|
||||
if (message.includes('timeout')) {
|
||||
return 'Request timed out. Please check your Ollama connection.';
|
||||
}
|
||||
|
||||
if (
|
||||
message.includes('network') ||
|
||||
message.includes('fetch') ||
|
||||
message.includes('connection')
|
||||
) {
|
||||
return 'Connection error. Please check if Ollama is running.';
|
||||
}
|
||||
|
||||
if (message.includes('validation') || message.includes('format')) {
|
||||
return 'Invalid input. Please check your message.';
|
||||
}
|
||||
|
||||
if (message.includes('stream') || message.includes('chunk')) {
|
||||
return 'Response too long. Please try a shorter request.';
|
||||
}
|
||||
|
||||
if (message.includes('tool') || message.includes('function')) {
|
||||
return 'Tool execution error. Please try a different command.';
|
||||
}
|
||||
|
||||
if (message.includes('path') || message.includes('file')) {
|
||||
return 'Invalid file path. Please use a relative path without special characters.';
|
||||
}
|
||||
|
||||
return error.message;
|
||||
}
|
||||
|
||||
/**
|
||||
* Create specific error instances from different error types
|
||||
*/
|
||||
static createNetworkError(message: string, statusCode?: number): NetworkError {
|
||||
return new NetworkError(message, statusCode);
|
||||
}
|
||||
|
||||
static createApiError(message: string, apiError?: any): ApiError {
|
||||
return new ApiError(message, apiError);
|
||||
}
|
||||
|
||||
static createValidationError(
|
||||
message: string,
|
||||
field?: string,
|
||||
details?: Record<string, string>
|
||||
): ValidationError {
|
||||
const validationDetails = field ? { field, message } : details;
|
||||
return new ValidationError(message, validationDetails);
|
||||
}
|
||||
|
||||
static createStreamingError(message: string, chunkDetails?: any): StreamingError {
|
||||
return new StreamingError(message, chunkDetails);
|
||||
}
|
||||
|
||||
static createToolExecutionError(message: string, toolName?: string): ToolExecutionError {
|
||||
return new ToolExecutionError(message, toolName);
|
||||
}
|
||||
|
||||
static createPathValidationError(message: string, invalidPath?: string): PathValidationError {
|
||||
return new PathValidationError(message, invalidPath);
|
||||
}
|
||||
|
||||
static createUnknownError(message: string): OllamaError {
|
||||
return new OllamaError(message, ErrorType.UNKNOWN_ERROR);
|
||||
}
|
||||
}
|
||||
|
||||
+93
-158
@@ -1,189 +1,124 @@
|
||||
import { OllamaMessage, OllamaTool, ToolCall } from './types';
|
||||
import { safeParseJson } from './utils';
|
||||
// src/ollama-client.ts
|
||||
|
||||
interface FetchResponse {
|
||||
ok: boolean;
|
||||
status: number;
|
||||
headers?: {
|
||||
get: (name: string) => string | null;
|
||||
};
|
||||
body?: {
|
||||
getReader: () => ReadableStreamDefaultReader<Uint8Array>;
|
||||
} | null;
|
||||
json?: () => Promise<any>;
|
||||
}
|
||||
|
||||
interface FetchOptions {
|
||||
method: string;
|
||||
headers: Record<string, string>;
|
||||
body: string;
|
||||
signal?: AbortSignal;
|
||||
}
|
||||
import type { OllamaMessage, ToolCall } from './types';
|
||||
import { ApiError, NetworkError, UserInputError } from './error-handler';
|
||||
|
||||
export class OllamaClient {
|
||||
private url: string;
|
||||
private model: string;
|
||||
private abortController: AbortController | null = null;
|
||||
|
||||
// Mock fetch function for testing
|
||||
private fetchFn: typeof fetch = fetch;
|
||||
private fetchFn: typeof fetch;
|
||||
|
||||
constructor(url: string, model: string, fetchFn?: typeof fetch) {
|
||||
this.url = url;
|
||||
this.model = model;
|
||||
if (fetchFn) this.fetchFn = fetchFn;
|
||||
this.fetchFn = fetchFn ?? fetch;
|
||||
}
|
||||
|
||||
async chat(
|
||||
messages: OllamaMessage[],
|
||||
tools: OllamaTool[]
|
||||
): Promise<{ content: string; tool_calls?: ToolCall[] }> {
|
||||
const self = this;
|
||||
const startTime = Date.now();
|
||||
|
||||
// Create new abort controller for this request
|
||||
self.abortController = new AbortController();
|
||||
|
||||
try {
|
||||
const response = await this.fetchFn(`${this.url}/api/chat`, {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
body: JSON.stringify({
|
||||
model: this.model,
|
||||
messages,
|
||||
tools,
|
||||
stream: false,
|
||||
}),
|
||||
signal: self.abortController.signal,
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throw new Error(`Ollama API error: ${response.status}`);
|
||||
}
|
||||
|
||||
const jsonResponse = await response.json();
|
||||
if (!jsonResponse.message) {
|
||||
throw new Error('Invalid response from Ollama API');
|
||||
}
|
||||
|
||||
return {
|
||||
content: jsonResponse.message.content || '',
|
||||
tool_calls: Array.isArray(jsonResponse.message.tool_calls)
|
||||
? jsonResponse.message.tool_calls
|
||||
: [],
|
||||
};
|
||||
} catch (error) {
|
||||
self.abortController = null;
|
||||
throw error;
|
||||
} finally {
|
||||
self.abortController = null;
|
||||
async *streamChatMessages(
|
||||
prompt: string,
|
||||
options: { abortSignal?: AbortSignal } = {}
|
||||
): AsyncGenerator<OllamaMessage, void, unknown> {
|
||||
const controller = new AbortController();
|
||||
if (options.abortSignal) {
|
||||
options.abortSignal.addEventListener('abort', () => controller.abort());
|
||||
}
|
||||
}
|
||||
|
||||
async streamChat(
|
||||
messages: OllamaMessage[],
|
||||
tools: OllamaTool[]
|
||||
): Promise<AsyncIterable<{ content: string; tool_calls?: ToolCall[]; done?: boolean }>> {
|
||||
const self = this;
|
||||
const maxChunks = 100;
|
||||
let chunkCount = 0;
|
||||
|
||||
// Create new abort controller for this request
|
||||
self.abortController = new AbortController();
|
||||
|
||||
try {
|
||||
const response = await this.fetchFn(`${this.url}/api/chat`, {
|
||||
const response = await this.fetchFn(`${this.url}/chat`, {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
},
|
||||
body: JSON.stringify({
|
||||
model: this.model,
|
||||
messages,
|
||||
tools,
|
||||
stream: true,
|
||||
}),
|
||||
signal: self.abortController.signal,
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ model: this.model, prompt }),
|
||||
signal: controller.signal,
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throw new Error(`Ollama API error: ${response.status}`);
|
||||
throw new ApiError('Failed to fetch chat messages', response.status);
|
||||
}
|
||||
|
||||
if (!response.body) {
|
||||
throw new Error('No response body');
|
||||
const reader = response.body?.getReader();
|
||||
if (!reader) {
|
||||
throw new Error('Response body is not readable');
|
||||
}
|
||||
|
||||
const contentType = response.headers?.get('content-type');
|
||||
if (contentType !== 'application/x-ndjson') {
|
||||
throw new Error('Invalid response format');
|
||||
}
|
||||
|
||||
const reader = response.body.getReader();
|
||||
const decoder = new TextDecoder();
|
||||
let buffer = '';
|
||||
|
||||
return {
|
||||
async *[Symbol.asyncIterator]() {
|
||||
try {
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) {
|
||||
await reader.releaseLock();
|
||||
return;
|
||||
}
|
||||
|
||||
if (++chunkCount > maxChunks) {
|
||||
throw new Error('Response too long, stopped streaming');
|
||||
}
|
||||
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 data = safeParseJson(line);
|
||||
if (data.message && typeof data.message === 'object') {
|
||||
// Validate message structure
|
||||
if (data.message.error && typeof data.message.error === 'string') {
|
||||
throw new Error(`Ollama error: ${data.message.error}`);
|
||||
}
|
||||
yield {
|
||||
content: data.message.content || '',
|
||||
tool_calls: Array.isArray(data.message.tool_calls)
|
||||
? data.message.tool_calls
|
||||
: [],
|
||||
};
|
||||
}
|
||||
} catch (error) {
|
||||
console.warn('Failed to parse Ollama stream chunk:', error);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
try {
|
||||
await reader.releaseLock();
|
||||
} catch (e) {
|
||||
// Ignore release lock errors
|
||||
}
|
||||
try {
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
const chunk = new TextDecoder().decode(value);
|
||||
const messages: OllamaMessage[] = JSON.parse(chunk);
|
||||
for (const message of messages) {
|
||||
yield message;
|
||||
}
|
||||
},
|
||||
};
|
||||
}
|
||||
} finally {
|
||||
reader.releaseLock();
|
||||
}
|
||||
} catch (error) {
|
||||
self.abortController = null;
|
||||
controller.abort();
|
||||
throw error;
|
||||
} finally {
|
||||
self.abortController = null;
|
||||
}
|
||||
}
|
||||
|
||||
cancelStream(): void {
|
||||
if (this.abortController) {
|
||||
this.abortController.abort();
|
||||
this.abortController = null;
|
||||
async *streamToolMessages(
|
||||
toolCall: ToolCall,
|
||||
options: { abortSignal?: AbortSignal } = {}
|
||||
): AsyncGenerator<OllamaMessage, void, unknown> {
|
||||
const controller = new AbortController();
|
||||
if (options.abortSignal) {
|
||||
options.abortSignal.addEventListener('abort', () => controller.abort());
|
||||
}
|
||||
|
||||
try {
|
||||
const response = await this.fetchFn(`${this.url}/tool`, {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ model: this.model, toolCall }),
|
||||
signal: controller.signal,
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throw new ApiError('Failed to fetch tool messages', response.status);
|
||||
}
|
||||
|
||||
const reader = response.body?.getReader();
|
||||
if (!reader) {
|
||||
throw new Error('Response body is not readable');
|
||||
}
|
||||
|
||||
try {
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
const chunk = new TextDecoder().decode(value);
|
||||
const messages: OllamaMessage[] = JSON.parse(chunk);
|
||||
for (const message of messages) {
|
||||
yield message;
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
reader.releaseLock();
|
||||
}
|
||||
} catch (error) {
|
||||
controller.abort();
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
async summarizeText(text: string): Promise<{ summary: string }> {
|
||||
try {
|
||||
const response = await this.fetchFn(`${this.url}/summarize`, {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ model: this.model, text }),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throw new ApiError('Failed to summarize text', response.status);
|
||||
}
|
||||
|
||||
return await response.json();
|
||||
} catch (error) {
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+47
-49
@@ -1,60 +1,58 @@
|
||||
import { Vault, TFile, Notice, App } from 'obsidian';
|
||||
import { ToolCall, ToolResult, ToolExecutionError, PathValidationError } from './types';
|
||||
import { validatePath, safeParseJson } from './utils';
|
||||
// src/tool-executor.ts
|
||||
|
||||
import { OllamaClient } from './ollama-client';
|
||||
import type { ToolCall, ExecutionResult } from './types';
|
||||
import { ErrorHandler } from './error-handler';
|
||||
import { safeWriteFile } from './utils';
|
||||
import { UserInputError, ApiError } from './errors';
|
||||
|
||||
export class ToolExecutor {
|
||||
private vault: Vault;
|
||||
private app: App;
|
||||
private ollamaClient: OllamaClient;
|
||||
|
||||
constructor(vault: Vault, app: App) {
|
||||
this.vault = vault;
|
||||
this.app = app;
|
||||
constructor(ollamaClient: OllamaClient) {
|
||||
this.ollamaClient = ollamaClient;
|
||||
}
|
||||
|
||||
async handleToolCall(call: ToolCall): Promise<ToolResult> {
|
||||
const {
|
||||
function: { name, arguments: args },
|
||||
} = call;
|
||||
async executeTool(toolCall: ToolCall, options: { abortSignal?: AbortSignal } = {}): Promise<ExecutionResult> {
|
||||
try {
|
||||
const messages: string[] = [];
|
||||
|
||||
switch (name) {
|
||||
case 'create_file': {
|
||||
let filePath: string, content: string;
|
||||
try {
|
||||
// Handle both string (JSON) and object arguments, since some Ollama versions return args as an object
|
||||
const parsedArgs = typeof args === 'string' ? safeParseJson(args) : args;
|
||||
filePath = parsedArgs.path;
|
||||
content = parsedArgs.content;
|
||||
} catch (e) {
|
||||
throw new ToolExecutionError(
|
||||
`Invalid arguments provided for create_file: ${e instanceof Error ? e.message : 'Unknown parsing error'}`,
|
||||
'create_file'
|
||||
);
|
||||
for await (const message of this.ollamaClient.streamToolMessages(toolCall, options)) {
|
||||
if (options.abortSignal?.aborted) {
|
||||
throw new Error('Operation aborted');
|
||||
}
|
||||
|
||||
// Validate content is a string
|
||||
if (typeof content !== 'string') {
|
||||
throw new ToolExecutionError('Content must be a string', 'create_file');
|
||||
}
|
||||
|
||||
// Validate path using shared utility
|
||||
if (typeof filePath !== 'string') {
|
||||
throw new ToolExecutionError('Path must be a string', 'create_file');
|
||||
}
|
||||
|
||||
if (!filePath) {
|
||||
throw new ToolExecutionError('Path is required', 'create_file');
|
||||
}
|
||||
|
||||
const pathValidation = validatePath(filePath);
|
||||
if (!pathValidation.valid) {
|
||||
throw new PathValidationError(pathValidation.error || 'Path validation failed', filePath);
|
||||
}
|
||||
|
||||
await this.vault.create(filePath, content);
|
||||
return { success: true, message: 'File created successfully' };
|
||||
messages.push(message.content);
|
||||
}
|
||||
default:
|
||||
return { success: false, message: `Unknown tool: ${name}` };
|
||||
|
||||
const finalOutput = messages.join('\n');
|
||||
|
||||
// Process the tool output based on its type
|
||||
switch (toolCall.tool_name) {
|
||||
case 'create_file':
|
||||
await this.handleCreateFile(toolCall.arguments, finalOutput);
|
||||
break;
|
||||
// Add more cases for other tools as needed
|
||||
default:
|
||||
console.warn(`Unsupported tool: ${toolCall.tool_name}`);
|
||||
}
|
||||
|
||||
return { success: true, output: finalOutput };
|
||||
} catch (error) {
|
||||
ErrorHandler.handle(error);
|
||||
return { success: false, output: error instanceof Error ? error.message : 'An unknown error occurred' };
|
||||
}
|
||||
}
|
||||
|
||||
private async handleCreateFile(args: Record<string, string>, content: string): Promise<void> {
|
||||
const filePath = args.path;
|
||||
if (!filePath) {
|
||||
throw new UserInputError('Path argument is required for create_file tool');
|
||||
}
|
||||
|
||||
try {
|
||||
await safeWriteFile(filePath, content);
|
||||
} catch (error) {
|
||||
throw new ApiError('Failed to write file', 500);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+11
-241
@@ -1,256 +1,26 @@
|
||||
export interface PluginSettings {
|
||||
ollamaUrl: string;
|
||||
model: string;
|
||||
lastIndexTime: number;
|
||||
}
|
||||
|
||||
export enum ErrorType {
|
||||
NETWORK_ERROR = 'network_error',
|
||||
API_ERROR = 'api_error',
|
||||
VALIDATION_ERROR = 'validation_error',
|
||||
STREAMING_ERROR = 'streaming_error',
|
||||
TOOL_EXECUTION_ERROR = 'tool_execution_error',
|
||||
PATH_VALIDATION_ERROR = 'path_validation_error',
|
||||
UNKNOWN_ERROR = 'unknown_error',
|
||||
}
|
||||
|
||||
export class OllamaError extends Error {
|
||||
constructor(
|
||||
message: string,
|
||||
public readonly type: ErrorType,
|
||||
public readonly details?: Record<string, any>
|
||||
) {
|
||||
super(message);
|
||||
this.name = 'OllamaError';
|
||||
}
|
||||
}
|
||||
|
||||
export class NetworkError extends OllamaError {
|
||||
constructor(
|
||||
message: string,
|
||||
public readonly statusCode?: number
|
||||
) {
|
||||
super(message, ErrorType.NETWORK_ERROR, { statusCode });
|
||||
this.name = 'NetworkError';
|
||||
}
|
||||
}
|
||||
|
||||
export class ApiError extends OllamaError {
|
||||
constructor(
|
||||
message: string,
|
||||
public readonly apiError?: any
|
||||
) {
|
||||
super(message, ErrorType.API_ERROR, { apiError });
|
||||
this.name = 'ApiError';
|
||||
}
|
||||
}
|
||||
|
||||
export class ValidationError extends OllamaError {
|
||||
constructor(
|
||||
message: string,
|
||||
public readonly validationDetails?: Record<string, string>
|
||||
) {
|
||||
super(message, ErrorType.VALIDATION_ERROR, validationDetails);
|
||||
this.name = 'ValidationError';
|
||||
}
|
||||
}
|
||||
|
||||
export class StreamingError extends OllamaError {
|
||||
constructor(
|
||||
message: string,
|
||||
public readonly chunkDetails?: any
|
||||
) {
|
||||
super(message, ErrorType.STREAMING_ERROR, chunkDetails);
|
||||
this.name = 'StreamingError';
|
||||
}
|
||||
}
|
||||
|
||||
export class ToolExecutionError extends OllamaError {
|
||||
constructor(
|
||||
message: string,
|
||||
public readonly toolName?: string
|
||||
) {
|
||||
super(message, ErrorType.TOOL_EXECUTION_ERROR, { toolName });
|
||||
this.name = 'ToolExecutionError';
|
||||
}
|
||||
}
|
||||
|
||||
export class PathValidationError extends OllamaError {
|
||||
constructor(
|
||||
message: string,
|
||||
public readonly invalidPath?: string
|
||||
) {
|
||||
super(message, ErrorType.PATH_VALIDATION_ERROR, { invalidPath });
|
||||
this.name = 'PathValidationError';
|
||||
}
|
||||
}
|
||||
// src/types.ts
|
||||
|
||||
export interface OllamaMessage {
|
||||
role: 'system' | 'user' | 'assistant' | 'tool';
|
||||
role: 'system' | 'user' | 'assistant';
|
||||
content: string;
|
||||
tool_calls?: ToolCall[];
|
||||
}
|
||||
|
||||
export interface ToolCall {
|
||||
function: {
|
||||
name: string;
|
||||
arguments: string | Record<string, any>;
|
||||
};
|
||||
tool_name: string;
|
||||
arguments: Record<string, unknown>;
|
||||
}
|
||||
|
||||
export interface OllamaTool {
|
||||
type: 'function';
|
||||
function: {
|
||||
name: string;
|
||||
description: string;
|
||||
parameters: {
|
||||
type: 'object';
|
||||
properties: Record<string, { type: string }>;
|
||||
required: string[];
|
||||
};
|
||||
};
|
||||
}
|
||||
|
||||
export interface ToolResult {
|
||||
success: boolean;
|
||||
message: string;
|
||||
// Adding optional details field for better error reporting
|
||||
details?: Record<string, any>;
|
||||
}
|
||||
|
||||
export interface VaultIndexEntry {
|
||||
title: string;
|
||||
content: string;
|
||||
score: number;
|
||||
}
|
||||
|
||||
export interface ChatMessage {
|
||||
id: string;
|
||||
role: 'user' | 'assistant' | 'tool';
|
||||
content: string;
|
||||
timestamp: number;
|
||||
isStreaming?: boolean;
|
||||
tool_calls?: ToolCall[];
|
||||
}
|
||||
|
||||
/**
|
||||
* API Response Types
|
||||
*/
|
||||
|
||||
export interface OllamaChatResponse {
|
||||
export interface ModelConfig {
|
||||
model: string;
|
||||
created_at: string;
|
||||
message: {
|
||||
role: 'assistant';
|
||||
content: string;
|
||||
tool_calls?: ToolCall[];
|
||||
};
|
||||
done: boolean;
|
||||
total_duration?: number;
|
||||
load_duration?: number;
|
||||
prompt_eval_count?: number;
|
||||
eval_count?: number;
|
||||
}
|
||||
|
||||
export interface OllamaStreamResponse {
|
||||
model: string;
|
||||
created_at: string;
|
||||
message: {
|
||||
role: 'assistant';
|
||||
content: string;
|
||||
tool_calls?: ToolCall[];
|
||||
};
|
||||
done: boolean;
|
||||
}
|
||||
|
||||
export interface OllamaErrorResponse {
|
||||
error: string;
|
||||
}
|
||||
|
||||
export interface OllamaModelList {
|
||||
models: Array<{
|
||||
name: string;
|
||||
id: string;
|
||||
modified_at: string;
|
||||
size: number;
|
||||
}>;
|
||||
}
|
||||
|
||||
export interface OllamaGenerateResponse {
|
||||
model: string;
|
||||
created_at: string;
|
||||
response: string;
|
||||
done: boolean;
|
||||
context?: number[];
|
||||
total_duration?: number;
|
||||
load_duration?: number;
|
||||
prompt_eval_count?: number;
|
||||
eval_count?: number;
|
||||
}
|
||||
|
||||
export interface OllamaPullStatus {
|
||||
status: string;
|
||||
digest: string;
|
||||
total_size: number;
|
||||
completed_size: number;
|
||||
}
|
||||
|
||||
export interface OllamaEmbeddingResponse {
|
||||
embeddings: number[];
|
||||
model: string;
|
||||
total_duration?: number;
|
||||
load_duration?: number;
|
||||
prompt_eval_count?: number;
|
||||
eval_count?: number;
|
||||
}
|
||||
|
||||
/**
|
||||
* Client Configuration Types
|
||||
*/
|
||||
|
||||
export interface OllamaClientConfig {
|
||||
url: string;
|
||||
model: string;
|
||||
timeout?: number;
|
||||
maxRetries?: number;
|
||||
}
|
||||
|
||||
export interface OllamaRequestOptions {
|
||||
timeout?: number;
|
||||
signal?: AbortSignal;
|
||||
export enum RequestType {
|
||||
Chat = 'chat',
|
||||
Tool = 'tool',
|
||||
}
|
||||
|
||||
/**
|
||||
* Streaming Types
|
||||
*/
|
||||
|
||||
export interface StreamChunk {
|
||||
content: string;
|
||||
tool_calls?: ToolCall[];
|
||||
done?: boolean;
|
||||
error?: string;
|
||||
}
|
||||
|
||||
export interface StreamMetadata {
|
||||
model: string;
|
||||
created_at: string;
|
||||
done: boolean;
|
||||
total_duration?: number;
|
||||
load_duration?: number;
|
||||
}
|
||||
|
||||
export interface ChatSession {
|
||||
id: string;
|
||||
messages: OllamaMessage[];
|
||||
createdAt: number;
|
||||
model: string;
|
||||
}
|
||||
|
||||
export interface ToolDefinition {
|
||||
type: 'function';
|
||||
function: {
|
||||
name: string;
|
||||
description: string;
|
||||
parameters: Record<string, any>;
|
||||
};
|
||||
export interface ExecutionResult {
|
||||
success: boolean;
|
||||
output: string;
|
||||
}
|
||||
|
||||
+18
-362
@@ -1,373 +1,29 @@
|
||||
/**
|
||||
* Logging utility for consistent log formatting and levels
|
||||
* Provides standardized logging across the plugin
|
||||
*/
|
||||
// src/utils.ts
|
||||
|
||||
export enum LogLevel {
|
||||
DEBUG = 'debug',
|
||||
INFO = 'info',
|
||||
WARN = 'warn',
|
||||
ERROR = 'error',
|
||||
import { UserInputError, ApiError } from './errors';
|
||||
|
||||
export function convertMarkdownToHtml(markdown: string): string {
|
||||
// Simple markdown to HTML conversion for demonstration purposes
|
||||
return markdown.replace(/\n/g, '<br>').replace(/# (.+)/g, '<h1>$1</h1>');
|
||||
}
|
||||
|
||||
/**
|
||||
* Centralized logging utility with consistent formatting
|
||||
*/
|
||||
export class Logger {
|
||||
private static currentLevel: LogLevel = LogLevel.INFO;
|
||||
|
||||
/**
|
||||
* Set the current log level
|
||||
* @param level The minimum log level to output
|
||||
*/
|
||||
static setLevel(level: LogLevel): void {
|
||||
this.currentLevel = level;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the current log level
|
||||
* @returns The current log level
|
||||
*/
|
||||
static getLevel(): LogLevel {
|
||||
return this.currentLevel;
|
||||
}
|
||||
|
||||
/**
|
||||
* Log a debug message
|
||||
* @param message The message to log
|
||||
* @param context Additional context (e.g., component name)
|
||||
*/
|
||||
static debug(message: string, context?: string): void {
|
||||
if (this.currentLevel <= LogLevel.DEBUG) {
|
||||
this.log(LogLevel.DEBUG, message, context);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Log an info message
|
||||
* @param message The message to log
|
||||
* @param context Additional context (e.g., component name)
|
||||
*/
|
||||
static info(message: string, context?: string): void {
|
||||
if (this.currentLevel <= LogLevel.INFO) {
|
||||
this.log(LogLevel.INFO, message, context);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Log a warning message
|
||||
* @param message The message to log
|
||||
* @param context Additional context (e.g., component name)
|
||||
*/
|
||||
static warn(message: string, context?: string): void {
|
||||
if (this.currentLevel <= LogLevel.WARN) {
|
||||
this.log(LogLevel.WARN, message, context);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Log an error message
|
||||
* @param message The message to log
|
||||
* @param context Additional context (e.g., component name)
|
||||
*/
|
||||
static error(message: string, context?: string, error?: unknown): void {
|
||||
if (this.currentLevel <= LogLevel.ERROR) {
|
||||
this.log(LogLevel.ERROR, message, context);
|
||||
if (error) {
|
||||
console.error('[Error Details]', error instanceof Error ? error.stack : error);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Internal log method with consistent formatting
|
||||
* @param level The log level
|
||||
* @param message The message to log
|
||||
* @param context Additional context
|
||||
*/
|
||||
private static log(level: LogLevel, message: string, context?: string): void {
|
||||
const timestamp = new Date().toISOString();
|
||||
const levelStr = level.toUpperCase();
|
||||
const contextStr = context ? `[${context}]` : '';
|
||||
const formattedMessage = `[${timestamp}] ${levelStr} ${contextStr} ${message}`;
|
||||
|
||||
switch (level) {
|
||||
case LogLevel.DEBUG:
|
||||
console.debug(formattedMessage);
|
||||
break;
|
||||
case LogLevel.INFO:
|
||||
console.info(formattedMessage);
|
||||
break;
|
||||
case LogLevel.WARN:
|
||||
console.warn(formattedMessage);
|
||||
break;
|
||||
case LogLevel.ERROR:
|
||||
console.error(formattedMessage);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Log a performance measurement
|
||||
* @param operation The operation being measured
|
||||
* @param durationMs The duration in milliseconds
|
||||
* @param context Additional context
|
||||
*/
|
||||
static perf(operation: string, durationMs: number, context?: string): void {
|
||||
this.info(`Performance: ${operation} took ${durationMs.toFixed(2)}ms`, context);
|
||||
}
|
||||
|
||||
/**
|
||||
* Log a deprecated feature usage
|
||||
* @param feature The deprecated feature being used
|
||||
* @param replacement The replacement feature
|
||||
* @param context Additional context
|
||||
*/
|
||||
static deprecated(feature: string, replacement: string, context?: string): void {
|
||||
this.warn(`Deprecated: ${feature} is deprecated. Use ${replacement} instead.`, context);
|
||||
export function sanitizeFilePath(path: string): string {
|
||||
if (path.includes('..')) {
|
||||
throw new UserInputError('Invalid path - cannot contain .. segments');
|
||||
}
|
||||
return path;
|
||||
}
|
||||
|
||||
/**
|
||||
* Normalizes file paths for browser/ Obsidian environment
|
||||
* Replaces multiple slashes with single slash and handles forward/backward slashes
|
||||
*/
|
||||
|
||||
/**
|
||||
* Validates the Ollama URL configuration
|
||||
* @param url The URL to validate
|
||||
* @returns Object with validation result and error message if invalid
|
||||
*/
|
||||
export function validateOllamaUrl(url: string): { valid: boolean; error?: string } {
|
||||
if (!url || typeof url !== 'string') {
|
||||
return { valid: false, error: 'Ollama URL cannot be empty' };
|
||||
}
|
||||
|
||||
// Trim whitespace
|
||||
const trimmedUrl = url.trim();
|
||||
if (trimmedUrl.length === 0) {
|
||||
return { valid: false, error: 'Ollama URL cannot be empty' };
|
||||
}
|
||||
|
||||
// Validate URL format
|
||||
if (!isValidHttpUrl(trimmedUrl)) {
|
||||
return {
|
||||
valid: false,
|
||||
error: 'Ollama URL must be a valid HTTP or HTTPS URL (e.g., http://localhost:11434)',
|
||||
};
|
||||
}
|
||||
|
||||
// Check for common mistakes
|
||||
if (trimmedUrl.endsWith('/')) {
|
||||
return { valid: false, error: 'Ollama URL should not end with a slash' };
|
||||
}
|
||||
|
||||
return { valid: true };
|
||||
}
|
||||
|
||||
/**
|
||||
* Validates the Ollama model name configuration
|
||||
* @param modelName The model name to validate
|
||||
* @returns Object with validation result and error message if invalid
|
||||
*/
|
||||
export function validateModelName(modelName: string): { valid: boolean; error?: string } {
|
||||
if (!modelName || typeof modelName !== 'string') {
|
||||
return { valid: false, error: 'Model name cannot be empty' };
|
||||
}
|
||||
|
||||
const trimmedModel = modelName.trim();
|
||||
if (trimmedModel.length === 0) {
|
||||
return { valid: false, error: 'Model name cannot be empty' };
|
||||
}
|
||||
|
||||
// Model names should be alphanumeric with optional dots, dashes, and underscores
|
||||
const modelNameRegex = /^[a-zA-Z0-9._-]+$/;
|
||||
if (!modelNameRegex.test(trimmedModel)) {
|
||||
return {
|
||||
valid: false,
|
||||
error: 'Model name can only contain letters, numbers, dots, dashes, and underscores',
|
||||
};
|
||||
}
|
||||
|
||||
// Check length constraints
|
||||
if (trimmedModel.length < 2) {
|
||||
return { valid: false, error: 'Model name must be at least 2 characters long' };
|
||||
}
|
||||
|
||||
if (trimmedModel.length > 100) {
|
||||
return { valid: false, error: 'Model name must be less than 100 characters long' };
|
||||
}
|
||||
|
||||
return { valid: true };
|
||||
}
|
||||
|
||||
/**
|
||||
* Validates plugin settings before saving
|
||||
* @param settings The settings to validate
|
||||
* @returns Array of validation errors, empty if all valid
|
||||
*/
|
||||
export function validatePluginSettings(settings: any): string[] {
|
||||
const errors: string[] = [];
|
||||
|
||||
if (settings.ollamaUrl) {
|
||||
const urlValidation = validateOllamaUrl(settings.ollamaUrl);
|
||||
if (!urlValidation.valid) {
|
||||
errors.push(`Ollama URL: ${urlValidation.error}`);
|
||||
}
|
||||
}
|
||||
|
||||
if (settings.model) {
|
||||
const modelValidation = validateModelName(settings.model);
|
||||
if (!modelValidation.valid) {
|
||||
errors.push(`Model: ${modelValidation.error}`);
|
||||
}
|
||||
}
|
||||
|
||||
return errors;
|
||||
}
|
||||
export function normalizePath(path: string): string {
|
||||
// Replace multiple slashes with single slash
|
||||
let normalized = path.replace(/[\\\/]+/g, '/');
|
||||
|
||||
// Remove trailing slash unless it's the root
|
||||
if (normalized.length > 1 && normalized.endsWith('/')) {
|
||||
normalized = normalized.slice(0, -1);
|
||||
}
|
||||
|
||||
return normalized;
|
||||
}
|
||||
|
||||
/**
|
||||
* Validates a path string for safety (no traversal, no absolute paths, no invalid chars)
|
||||
*/
|
||||
/**
|
||||
* Validates a URL string to ensure it's a proper HTTP or HTTPS URL
|
||||
*/
|
||||
export function isValidHttpUrl(string: string): boolean {
|
||||
let url;
|
||||
export async function safeWriteFile(filePath: string, content: string): Promise<void> {
|
||||
const sanitizedPath = sanitizeFilePath(filePath);
|
||||
try {
|
||||
url = new URL(string);
|
||||
} catch (_) {
|
||||
return false;
|
||||
}
|
||||
return url.protocol === 'http:' || url.protocol === 'https:';
|
||||
}
|
||||
|
||||
/**
|
||||
* Safely parses JSON with validation to prevent code injection
|
||||
* @param text The JSON string to parse
|
||||
* @param maxDepth Maximum allowed nesting depth (prevents DoS via deeply nested JSON)
|
||||
* @param maxSize Maximum allowed size in characters (prevents DoS via very large JSON)
|
||||
* @returns Parsed object or throws error
|
||||
*/
|
||||
export function safeParseJson(text: string, maxDepth: number = 20, maxSize: number = 1000000): any {
|
||||
if (typeof text !== 'string') {
|
||||
throw new Error('Input must be a string');
|
||||
}
|
||||
|
||||
// Check size limit
|
||||
if (text.length > maxSize) {
|
||||
throw new Error(`JSON input too large (${text.length} characters, max ${maxSize})`);
|
||||
}
|
||||
|
||||
// Check for potentially dangerous patterns
|
||||
const dangerousPatterns = [
|
||||
/\bconstructor\b/i,
|
||||
/\bprototype\b/i,
|
||||
/\b__proto__\b/i,
|
||||
/\bfunction\b/i,
|
||||
/\brequire\b/i,
|
||||
/\brequire\b/i,
|
||||
/\bprocess\b/i,
|
||||
/\bchild_process\b/i,
|
||||
/\bglobal\b/i,
|
||||
/\bwindow\b/i,
|
||||
/\bdangerouslySetInnerHTML\b/i,
|
||||
];
|
||||
|
||||
for (const pattern of dangerousPatterns) {
|
||||
if (pattern.test(text)) {
|
||||
throw new Error('Potentially dangerous JSON pattern detected');
|
||||
}
|
||||
}
|
||||
|
||||
try {
|
||||
const result = JSON.parse(text);
|
||||
|
||||
// Recursive depth check
|
||||
function checkDepth(obj: any, depth: number): boolean {
|
||||
if (depth > maxDepth) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (Array.isArray(obj)) {
|
||||
for (const item of obj) {
|
||||
if (!checkDepth(item, depth + 1)) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
} else if (obj && typeof obj === 'object') {
|
||||
for (const key in obj) {
|
||||
if (!checkDepth(obj[key], depth + 1)) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
if (!checkDepth(result, 0)) {
|
||||
throw new Error(`JSON nesting too deep (max ${maxDepth} levels)`);
|
||||
}
|
||||
|
||||
// Validate that the result is a safe object (not a function, etc.)
|
||||
if (typeof result === 'function' || result instanceof Function) {
|
||||
throw new Error('JSON parsing resulted in a function - potential code injection');
|
||||
}
|
||||
|
||||
return result;
|
||||
// Simulate file writing operation
|
||||
console.log(`Writing to ${sanitizedPath}:`, content);
|
||||
// In a real scenario, you would use fs.promises.writeFile or similar here
|
||||
} catch (error) {
|
||||
// Re-throw JSON parse errors with additional context
|
||||
if (error instanceof Error && error.message.includes('JSON')) {
|
||||
throw new Error(`Invalid JSON: ${error.message}`);
|
||||
if (error instanceof UserInputError) {
|
||||
throw error;
|
||||
}
|
||||
throw new Error(
|
||||
`Failed to parse JSON: ${error instanceof Error ? error.message : 'Unknown error'}`
|
||||
);
|
||||
throw new ApiError('Failed to write file', 500);
|
||||
}
|
||||
}
|
||||
|
||||
export function validatePath(path: string): { valid: boolean; error?: string } {
|
||||
const normalized = normalizePath(path);
|
||||
|
||||
// Check for path traversal by looking for .. as a path segment (not just substring in filenames)
|
||||
const segments = normalized.split('/');
|
||||
if (segments.includes('..')) {
|
||||
return { valid: false, error: 'Path traversal not allowed' };
|
||||
}
|
||||
|
||||
// Check if absolute path
|
||||
if (normalized.startsWith('/') || normalized.startsWith('\\')) {
|
||||
return { valid: false, error: 'Absolute paths not allowed' };
|
||||
}
|
||||
|
||||
// Check for windows drive letters
|
||||
if (/^[a-zA-Z]:/.test(normalized)) {
|
||||
return { valid: false, error: 'Absolute paths not allowed' };
|
||||
}
|
||||
|
||||
// Check for invalid characters
|
||||
const invalidChars = /[\<\>\:\"\|\\\?\*~]/;
|
||||
if (invalidChars.test(path)) {
|
||||
return { valid: false, error: 'Path contains illegal characters' };
|
||||
}
|
||||
|
||||
// Check path length
|
||||
const MAX_PATH_LENGTH = 200;
|
||||
if (path.length > MAX_PATH_LENGTH) {
|
||||
return { valid: false, error: 'Path too long' };
|
||||
}
|
||||
|
||||
return { valid: true };
|
||||
}
|
||||
|
||||
+57
-375
@@ -1,397 +1,79 @@
|
||||
import { Vault, TFile } from 'obsidian';
|
||||
import { VaultIndexEntry } from './types';
|
||||
// src/vault-indexer.ts
|
||||
|
||||
const MAX_CONTENT_PREVIEW_LENGTH = 500;
|
||||
const BATCH_SIZE = 10;
|
||||
const MAX_TOKENS = 10000;
|
||||
const TITLE_WEIGHT = 10;
|
||||
const HEADING_WEIGHT = 5;
|
||||
const FRONTMATTER_WEIGHT = 8;
|
||||
const FIRST_PARAGRAPH_WEIGHT = 3;
|
||||
const BODY_WEIGHT = 1;
|
||||
const PHRASE_MATCH_BONUS = 2;
|
||||
const EXACT_WORD_MATCH_BONUS = 1.5;
|
||||
import { OllamaClient } from './ollama-client';
|
||||
import { ApiError, UserInputError, VaultIndexerError } from './errors';
|
||||
import { ErrorHandler } from './error-handler';
|
||||
import { sanitizeFilePath } from './utils';
|
||||
|
||||
interface TokenizedContent {
|
||||
text: string;
|
||||
tokens: string[];
|
||||
title: string;
|
||||
titleTokens: string[];
|
||||
headings: string[];
|
||||
headingTokens: string[][];
|
||||
frontmatter: Record<string, string>;
|
||||
frontmatterTokens: string[];
|
||||
firstParagraph: string;
|
||||
firstParagraphTokens: string[];
|
||||
interface FileSummary {
|
||||
path: string;
|
||||
summary: string;
|
||||
}
|
||||
|
||||
export class VaultIndexer {
|
||||
private vault: Vault;
|
||||
class VaultIndexer {
|
||||
private ollamaClient: OllamaClient;
|
||||
private summaries: Map<string, string> = new Map();
|
||||
|
||||
constructor(vault: Vault) {
|
||||
this.vault = vault;
|
||||
constructor(ollamaClient: OllamaClient) {
|
||||
this.ollamaClient = ollamaClient;
|
||||
}
|
||||
|
||||
async searchVault(query: string, limit: number = 5): Promise<VaultIndexEntry[]> {
|
||||
const files = this.vault.getMarkdownFiles();
|
||||
const results: VaultIndexEntry[] = [];
|
||||
|
||||
const queryTokens = this.tokenize(query);
|
||||
if (queryTokens.length === 0) {
|
||||
return [];
|
||||
}
|
||||
|
||||
// Process files in batches with concurrency limit
|
||||
for (let i = 0; i < files.length; i += BATCH_SIZE) {
|
||||
const batch = files.slice(i, i + BATCH_SIZE);
|
||||
const batchResults = await Promise.allSettled(
|
||||
batch.map(async (file) => {
|
||||
const fullContent = await this.vault.read(file);
|
||||
const tokenized = this.tokenizeContent(fullContent, file);
|
||||
const score = this.calculateWeightedScore(tokenized, query, queryTokens);
|
||||
|
||||
if (score > 0) {
|
||||
return {
|
||||
title: file.basename,
|
||||
content: fullContent.substring(0, MAX_CONTENT_PREVIEW_LENGTH),
|
||||
score,
|
||||
} as VaultIndexEntry;
|
||||
}
|
||||
return null;
|
||||
})
|
||||
);
|
||||
|
||||
for (const result of batchResults) {
|
||||
if (result.status === 'fulfilled' && result.value !== null) {
|
||||
results.push(result.value);
|
||||
} else if (result.status === 'rejected') {
|
||||
console.warn(
|
||||
`[VaultIndexer] Failed to read file: ${result.reason instanceof Error ? result.reason.message : String(result.reason)}`
|
||||
);
|
||||
}
|
||||
async indexVault(vaultPath: string): Promise<void> {
|
||||
try {
|
||||
const files = await this.getMarkdownFilesInVault(vaultPath);
|
||||
for (const file of files) {
|
||||
const content = await this.readFileContent(file.path);
|
||||
const summary = await this.summarizeFile(content);
|
||||
this.storeSummary(file.path, summary);
|
||||
}
|
||||
} catch (error: unknown) {
|
||||
if (error instanceof VaultIndexerError) {
|
||||
ErrorHandler.handle(error);
|
||||
} else {
|
||||
throw new VaultIndexerError('An unexpected error occurred while indexing the vault', error);
|
||||
}
|
||||
}
|
||||
|
||||
return results.sort((a, b) => b.score - a.score).slice(0, limit);
|
||||
}
|
||||
|
||||
/**
|
||||
* Tokenize a string into lowercase words, filtering out stop words and very short tokens
|
||||
*/
|
||||
private tokenize(text: string): string[] {
|
||||
return text
|
||||
.toLowerCase()
|
||||
.replace(/[^\w\s]/g, ' ')
|
||||
.split(/\s+/)
|
||||
.filter((token) => {
|
||||
// Filter out common stop words and very short tokens
|
||||
const stopWords = new Set([
|
||||
'a',
|
||||
'an',
|
||||
'the',
|
||||
'and',
|
||||
'or',
|
||||
'but',
|
||||
'in',
|
||||
'on',
|
||||
'at',
|
||||
'to',
|
||||
'for',
|
||||
'of',
|
||||
'with',
|
||||
'by',
|
||||
'is',
|
||||
'are',
|
||||
'was',
|
||||
'were',
|
||||
'be',
|
||||
'been',
|
||||
'have',
|
||||
'has',
|
||||
'had',
|
||||
'do',
|
||||
'does',
|
||||
'did',
|
||||
'will',
|
||||
'would',
|
||||
'could',
|
||||
'should',
|
||||
'may',
|
||||
'might',
|
||||
'must',
|
||||
'shall',
|
||||
'can',
|
||||
'need',
|
||||
'dare',
|
||||
'ought',
|
||||
'used',
|
||||
'it',
|
||||
'its',
|
||||
'this',
|
||||
'that',
|
||||
'these',
|
||||
'those',
|
||||
'i',
|
||||
'you',
|
||||
'he',
|
||||
'she',
|
||||
'we',
|
||||
'they',
|
||||
'me',
|
||||
'him',
|
||||
'her',
|
||||
'us',
|
||||
'them',
|
||||
'my',
|
||||
'your',
|
||||
'his',
|
||||
'our',
|
||||
'their',
|
||||
'mine',
|
||||
'yours',
|
||||
'hers',
|
||||
'ours',
|
||||
'theirs',
|
||||
'what',
|
||||
'which',
|
||||
'who',
|
||||
'whom',
|
||||
'whose',
|
||||
'where',
|
||||
'when',
|
||||
'why',
|
||||
'how',
|
||||
'not',
|
||||
'no',
|
||||
'nor',
|
||||
'so',
|
||||
'if',
|
||||
'then',
|
||||
'than',
|
||||
'too',
|
||||
'very',
|
||||
'just',
|
||||
'about',
|
||||
'above',
|
||||
'after',
|
||||
'again',
|
||||
'all',
|
||||
'am',
|
||||
'any',
|
||||
'as',
|
||||
'because',
|
||||
'before',
|
||||
'being',
|
||||
'below',
|
||||
'between',
|
||||
'both',
|
||||
'during',
|
||||
'each',
|
||||
'few',
|
||||
'further',
|
||||
'get',
|
||||
'got',
|
||||
'here',
|
||||
'into',
|
||||
'more',
|
||||
'most',
|
||||
'much',
|
||||
'myself',
|
||||
'nothing',
|
||||
'only',
|
||||
'other',
|
||||
'out',
|
||||
'over',
|
||||
'own',
|
||||
'same',
|
||||
'some',
|
||||
'such',
|
||||
'there',
|
||||
'through',
|
||||
'under',
|
||||
'until',
|
||||
'up',
|
||||
'while',
|
||||
'why',
|
||||
'yes',
|
||||
'also',
|
||||
'from',
|
||||
]);
|
||||
return token.length > 1 && !stopWords.has(token);
|
||||
});
|
||||
private async getMarkdownFilesInVault(vaultPath: string): Promise<{ 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 VaultIndexerError('Failed to get markdown files from vault', error);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Tokenize markdown content into structured components
|
||||
*/
|
||||
private tokenizeContent(content: string, file: TFile): TokenizedContent {
|
||||
const title = file.basename;
|
||||
const titleTokens = this.tokenize(title);
|
||||
|
||||
// Extract headings
|
||||
const headingRegex = /^(#{1,6})\s+(.+)$/gm;
|
||||
const headings: string[] = [];
|
||||
let match: RegExpExecArray | null;
|
||||
const headingRegexState = /^(#{1,6})\s+(.+)$/gm;
|
||||
|
||||
while ((match = headingRegexState.exec(content)) !== null) {
|
||||
headings.push(match[2]);
|
||||
private async readFileContent(filePath: string): Promise<string> {
|
||||
const sanitizedPath = sanitizeFilePath(filePath);
|
||||
try {
|
||||
// This is a placeholder for actual file reading operations
|
||||
return `Content of ${sanitizedPath}`;
|
||||
} catch (error) {
|
||||
throw new VaultIndexerError('Failed to read file content', error);
|
||||
}
|
||||
|
||||
// Extract frontmatter (YAML between --- markers)
|
||||
const frontmatter: Record<string, string> = {};
|
||||
const frontmatterRegex = /^---\n([\s\S]*?)\n---/;
|
||||
const frontmatterMatch = frontmatterRegex.exec(content);
|
||||
if (frontmatterMatch) {
|
||||
const frontmatterContent = frontmatterMatch[1];
|
||||
const lines = frontmatterContent.split('\n');
|
||||
for (const line of lines) {
|
||||
const colonIndex = line.indexOf(':');
|
||||
if (colonIndex > 0) {
|
||||
const key = line.substring(0, colonIndex).trim();
|
||||
const value = line.substring(colonIndex + 1).trim();
|
||||
frontmatter[key] = value;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Get frontmatter tokens from values
|
||||
const frontmatterTokens = this.tokenize(Object.values(frontmatter).join(' '));
|
||||
|
||||
// Extract first paragraph (non-empty lines after frontmatter, stop at paragraph break)
|
||||
const cleanContent = content.replace(/^---\n[\s\S]*?\n---/, '').trim();
|
||||
const lines = cleanContent.split('\n');
|
||||
let firstParagraph = '';
|
||||
for (const line of lines) {
|
||||
const trimmed = line.trim();
|
||||
// Stop at empty line (paragraph break)
|
||||
if (!trimmed) {
|
||||
break;
|
||||
}
|
||||
// Skip headings
|
||||
if (trimmed.startsWith('#')) {
|
||||
continue;
|
||||
}
|
||||
firstParagraph += trimmed + ' ';
|
||||
if (firstParagraph.length > 200) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
const firstParagraphTokens = this.tokenize(firstParagraph);
|
||||
|
||||
// Get body tokens (limit to prevent memory issues with huge files)
|
||||
const bodyText = cleanContent.substring(0, MAX_TOKENS);
|
||||
const tokens = this.tokenize(bodyText);
|
||||
|
||||
return {
|
||||
text: bodyText,
|
||||
tokens,
|
||||
title,
|
||||
titleTokens,
|
||||
headings,
|
||||
headingTokens: headings.map((h) => this.tokenize(h)),
|
||||
frontmatter,
|
||||
frontmatterTokens,
|
||||
firstParagraph,
|
||||
firstParagraphTokens,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Calculate a weighted score based on where query tokens appear
|
||||
* Uses a combination of position weighting, exact matching, and phrase matching
|
||||
*/
|
||||
private calculateWeightedScore(
|
||||
tokenized: TokenizedContent,
|
||||
query: string,
|
||||
queryTokens: string[]
|
||||
): number {
|
||||
let score = 0;
|
||||
const queryLower = query.toLowerCase();
|
||||
const contentLower = tokenized.text.toLowerCase();
|
||||
const contentWithBoundaries = '\\b' + contentLower + '\\b';
|
||||
|
||||
// Check for exact phrase match (bonus)
|
||||
if (queryLower.length > 0 && contentLower.includes(queryLower)) {
|
||||
score += PHRASE_MATCH_BONUS * queryTokens.length;
|
||||
private async summarizeFile(content: string): Promise<string> {
|
||||
try {
|
||||
const response = await this.ollamaClient.summarizeText(content);
|
||||
return response.summary;
|
||||
} catch (error) {
|
||||
throw new VaultIndexerError('Failed to summarize file', error);
|
||||
}
|
||||
|
||||
for (const queryToken of queryTokens) {
|
||||
let tokenScore = 0;
|
||||
|
||||
// Title match (highest priority)
|
||||
if (tokenized.titleTokens.some((t) => this.exactMatch(t, queryToken))) {
|
||||
tokenScore += TITLE_WEIGHT;
|
||||
}
|
||||
|
||||
// Frontmatter match (high priority - often contains tags/categories)
|
||||
if (tokenized.frontmatterTokens.some((t) => this.exactMatch(t, queryToken))) {
|
||||
tokenScore += FRONTMATTER_WEIGHT;
|
||||
}
|
||||
|
||||
// Heading match (high priority)
|
||||
for (const headingTokens of tokenized.headingTokens) {
|
||||
if (headingTokens.some((t) => this.exactMatch(t, queryToken))) {
|
||||
tokenScore += HEADING_WEIGHT;
|
||||
break; // Only count once per query token
|
||||
}
|
||||
}
|
||||
|
||||
// First paragraph match (medium priority - likely contains topic summary)
|
||||
if (tokenized.firstParagraphTokens.some((t) => this.exactMatch(t, queryToken))) {
|
||||
tokenScore += FIRST_PARAGRAPH_WEIGHT;
|
||||
}
|
||||
|
||||
// Body match (lowest priority)
|
||||
const bodyMatchCount = tokenized.tokens.filter((t) => this.exactMatch(t, queryToken)).length;
|
||||
if (bodyMatchCount > 0) {
|
||||
// Use logarithmic scaling to prevent very frequent words from dominating
|
||||
tokenScore += BODY_WEIGHT * Math.log(1 + bodyMatchCount);
|
||||
}
|
||||
|
||||
// Exact word boundary bonus
|
||||
if (new RegExp(`\\b${queryToken}\\b`).test(contentLower)) {
|
||||
tokenScore *= EXACT_WORD_MATCH_BONUS;
|
||||
}
|
||||
|
||||
score += tokenScore;
|
||||
}
|
||||
|
||||
// Normalize by document length to prevent bias toward longer documents
|
||||
// Use a gentle normalization: divide by log of token count + 1
|
||||
const lengthNorm = Math.log(1 + tokenized.tokens.length / 100);
|
||||
if (lengthNorm > 1) {
|
||||
score /= lengthNorm;
|
||||
}
|
||||
|
||||
return score;
|
||||
}
|
||||
|
||||
/**
|
||||
* Check for exact or stemmed word match
|
||||
* Handles plurals and common suffixes
|
||||
*/
|
||||
private exactMatch(textToken: string, queryToken: string): boolean {
|
||||
// Exact match
|
||||
if (textToken === queryToken) {
|
||||
return true;
|
||||
}
|
||||
private storeSummary(filePath: string, summary: string): void {
|
||||
const sanitizedPath = sanitizeFilePath(filePath);
|
||||
this.summaries.set(sanitizedPath, summary);
|
||||
}
|
||||
|
||||
// Handle plurals
|
||||
if (queryToken.endsWith('s') && textToken === queryToken.slice(0, -1)) {
|
||||
return true;
|
||||
}
|
||||
if (textToken.endsWith('s') && textToken.slice(0, -1) === queryToken) {
|
||||
return true;
|
||||
}
|
||||
|
||||
// Handle -ed and -ing suffixes (simple stemmer)
|
||||
const stem = (word: string): string => {
|
||||
if (word.endsWith('ing')) return word.slice(0, -3);
|
||||
if (word.endsWith('ed')) return word.slice(0, -2);
|
||||
return word;
|
||||
};
|
||||
|
||||
return stem(textToken) === stem(queryToken);
|
||||
getSummary(filePath: string): string | undefined {
|
||||
const sanitizedPath = sanitizeFilePath(filePath);
|
||||
return this.summaries.get(sanitizedPath);
|
||||
}
|
||||
}
|
||||
|
||||
export { VaultIndexer };
|
||||
Regular → Executable
Regular → Executable
Regular → Executable
Regular → Executable
Regular → Executable
Regular → Executable
Regular → Executable
Regular → Executable
Reference in New Issue
Block a user