From 36871bc0ac3bf567bd10111421a0ca6a894d5d06 Mon Sep 17 00:00:00 2001 From: Kwanghyun On Date: Sun, 27 Oct 2024 20:53:43 +0900 Subject: [PATCH] Handle llm setting errors during embedding generation --- src/components/chat-view/Chat.tsx | 23 ++---- src/utils/embedding.ts | 24 ++++++ src/utils/llm/exception.ts | 4 +- src/utils/llm/ollama.ts | 6 +- src/utils/openSettingsModal.ts | 12 +++ src/utils/vector-db/manager.ts | 121 ++++++++++++++++++------------ 6 files changed, 120 insertions(+), 70 deletions(-) create mode 100644 src/utils/openSettingsModal.ts diff --git a/src/components/chat-view/Chat.tsx b/src/components/chat-view/Chat.tsx index 6c6a7a0..c71f6a9 100644 --- a/src/components/chat-view/Chat.tsx +++ b/src/components/chat-view/Chat.tsx @@ -19,7 +19,6 @@ import { useLLM } from '../../contexts/llm-context' import { useRAG } from '../../contexts/rag-context' import { useSettings } from '../../contexts/settings-context' import { useChatHistory } from '../../hooks/useChatHistory' -import { OpenSettingsModal } from '../../OpenSettingsModal' import { ChatMessage, ChatUserMessage } from '../../types/chat' import { MentionableBlock, @@ -28,15 +27,16 @@ import { } from '../../types/mentionable' import { applyChangesToFile } from '../../utils/apply' import { - LLMABaseUrlNotSetException, LLMAPIKeyInvalidException, LLMAPIKeyNotSetException, + LLMBaseUrlNotSetException, } from '../../utils/llm/exception' import { getMentionableKey, serializeMentionable, } from '../../utils/mentionable' import { readTFileContent } from '../../utils/obsidian' +import { openSettingsModalWithError } from '../../utils/openSettingsModal' import { PromptGenerator } from '../../utils/promptGenerator' import ChatUserInput, { ChatUserInputRef } from './chat-input/ChatUserInput' @@ -261,14 +261,9 @@ const Chat = forwardRef((props, ref) => { if ( error instanceof LLMAPIKeyNotSetException || error instanceof LLMAPIKeyInvalidException || - error instanceof LLMABaseUrlNotSetException + error instanceof LLMBaseUrlNotSetException ) { - new OpenSettingsModal(app, error.message, () => { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const setting = (app as any).setting - setting.open() - setting.openTabById('smart-composer') - }).open() + openSettingsModalWithError(app, error.message) } else { new Notice(error.message) console.error('Failed to generate response', error) @@ -324,14 +319,10 @@ const Chat = forwardRef((props, ref) => { onError: (error) => { if ( error instanceof LLMAPIKeyNotSetException || - error instanceof LLMAPIKeyInvalidException + error instanceof LLMAPIKeyInvalidException || + error instanceof LLMBaseUrlNotSetException ) { - new OpenSettingsModal(app, error.message, () => { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const setting = (app as any).setting - setting.open() - setting.openTabById('smart-composer') - }).open() + openSettingsModalWithError(app, error.message) } else { new Notice(error.message) console.error('Failed to apply changes', error) diff --git a/src/utils/embedding.ts b/src/utils/embedding.ts index a34a031..55ddf85 100644 --- a/src/utils/embedding.ts +++ b/src/utils/embedding.ts @@ -2,6 +2,10 @@ import { OpenAI } from 'openai' import { EmbeddingModel } from '../types/embedding' +import { + LLMAPIKeyNotSetException, + LLMBaseUrlNotSetException, +} from './llm/exception' import { NoStainlessOpenAI } from './llm/ollama' export const getEmbeddingModel = ( @@ -21,6 +25,11 @@ export const getEmbeddingModel = ( name: 'text-embedding-3-small', dimension: 1536, getEmbedding: async (text: string) => { + if (!openai.apiKey) { + throw new LLMAPIKeyNotSetException( + 'OpenAI API key is missing. Please set it in settings menu.', + ) + } const embedding = await openai.embeddings.create({ model: 'text-embedding-3-small', input: text, @@ -38,6 +47,11 @@ export const getEmbeddingModel = ( name: 'text-embedding-3-large', dimension: 3072, getEmbedding: async (text: string) => { + if (!openai.apiKey) { + throw new LLMAPIKeyNotSetException( + 'OpenAI API key is missing. Please set it in settings menu.', + ) + } const embedding = await openai.embeddings.create({ model: 'text-embedding-3-large', input: text, @@ -56,6 +70,11 @@ export const getEmbeddingModel = ( name: 'nomic-embed-text', dimension: 768, getEmbedding: async (text: string) => { + if (!ollamaBaseUrl) { + throw new LLMBaseUrlNotSetException( + 'Ollama Address is missing. Please set it in settings menu.', + ) + } const embedding = await openai.embeddings.create({ model: 'nomic-embed-text', input: text, @@ -74,6 +93,11 @@ export const getEmbeddingModel = ( name: 'mxbai-embed-large', dimension: 1024, getEmbedding: async (text: string) => { + if (!ollamaBaseUrl) { + throw new LLMBaseUrlNotSetException( + 'Ollama Address is missing. Please set it in settings menu.', + ) + } const embedding = await openai.embeddings.create({ model: 'mxbai-embed-large', input: text, diff --git a/src/utils/llm/exception.ts b/src/utils/llm/exception.ts index 9dfdc69..2054b9d 100644 --- a/src/utils/llm/exception.ts +++ b/src/utils/llm/exception.ts @@ -12,9 +12,9 @@ export class LLMAPIKeyInvalidException extends Error { } } -export class LLMABaseUrlNotSetException extends Error { +export class LLMBaseUrlNotSetException extends Error { constructor(message: string) { super(message) - this.name = 'LLMABaseUrlNotSetException' + this.name = 'LLMBaseUrlNotSetException' } } diff --git a/src/utils/llm/ollama.ts b/src/utils/llm/ollama.ts index 498212f..20ab58a 100644 --- a/src/utils/llm/ollama.ts +++ b/src/utils/llm/ollama.ts @@ -11,7 +11,7 @@ import { } from 'src/types/llm/response' import { BaseLLMProvider } from './base' -import { LLMABaseUrlNotSetException } from './exception' +import { LLMBaseUrlNotSetException } from './exception' import { OpenAICompatibleProvider } from './openaiCompatibleProvider' export class NoStainlessOpenAI extends OpenAI { @@ -59,7 +59,7 @@ export class OllamaOpenAIProvider implements BaseLLMProvider { options?: LLMOptions, ): Promise { if (!this.ollamaBaseUrl) { - throw new LLMABaseUrlNotSetException( + throw new LLMBaseUrlNotSetException( 'Ollama Address is missing. Please set it in settings menu.', ) } @@ -70,7 +70,7 @@ export class OllamaOpenAIProvider implements BaseLLMProvider { options?: LLMOptions, ): Promise> { if (!this.ollamaBaseUrl) { - throw new LLMABaseUrlNotSetException( + throw new LLMBaseUrlNotSetException( 'Ollama Address is missing. Please set it in settings menu.', ) } diff --git a/src/utils/openSettingsModal.ts b/src/utils/openSettingsModal.ts new file mode 100644 index 0000000..c35e80d --- /dev/null +++ b/src/utils/openSettingsModal.ts @@ -0,0 +1,12 @@ +import { App } from 'obsidian' + +import { OpenSettingsModal } from '../OpenSettingsModal' + +export function openSettingsModalWithError(app: App, errorMessage: string) { + new OpenSettingsModal(app, errorMessage, () => { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const setting = (app as any).setting + setting.open() + setting.openTabById('smart-composer') + }).open() +} diff --git a/src/utils/vector-db/manager.ts b/src/utils/vector-db/manager.ts index dea5c35..43dc055 100644 --- a/src/utils/vector-db/manager.ts +++ b/src/utils/vector-db/manager.ts @@ -7,6 +7,12 @@ import { IndexProgress } from '../../components/chat-view/QueryProgress' import { PGLITE_DB_PATH } from '../../constants' import { InsertVector, SelectVector } from '../../db/schema' import { EmbeddingModel } from '../../types/embedding' +import { + LLMAPIKeyInvalidException, + LLMAPIKeyNotSetException, + LLMBaseUrlNotSetException, +} from '../llm/exception' +import { openSettingsModalWithError } from '../openSettingsModal' import { VectorDbRepository } from './repository' @@ -121,64 +127,81 @@ export class VectorDbManager { const embeddingChunks: InsertVector[] = [] const batchSize = 100 const limit = pLimit(50) + const abortController = new AbortController() const tasks = contentChunks.map((chunk) => limit(async () => { - await backOff( - async () => { - const embedding = await embeddingModel.getEmbedding(chunk.content) - const embeddedChunk = { - path: chunk.path, - mtime: chunk.mtime, - content: chunk.content, - embedding, - metadata: chunk.metadata, - } - embeddingChunks.push(embeddedChunk) - embeddingProgress.completed++ - updateProgress?.({ - completedChunks: embeddingProgress.completed, - totalChunks: contentChunks.length, - totalFiles: filesToIndex.length, - }) - - // Insert vectors in batches - if ( - embeddingChunks.length >= - embeddingProgress.inserted + batchSize || - embeddingChunks.length === contentChunks.length - ) { - await this.repository.insertVectors( - embeddingChunks.slice( - embeddingProgress.inserted, - embeddingProgress.inserted + batchSize, - ), - embeddingModel, - ) - embeddingProgress.inserted += batchSize - } - }, - { - numOfAttempts: 5, - startingDelay: 1000, - timeMultiple: 1.5, - jitter: 'full', - retry: (error) => { - console.error(error) - const isRateLimitError = - error.status === 429 && - error.message.toLowerCase().includes('rate limit') - return !!isRateLimitError // retry only for rate limit errors + if (abortController.signal.aborted) { + throw new Error('Operation was aborted') + } + try { + await backOff( + async () => { + const embedding = await embeddingModel.getEmbedding(chunk.content) + const embeddedChunk = { + path: chunk.path, + mtime: chunk.mtime, + content: chunk.content, + embedding, + metadata: chunk.metadata, + } + embeddingChunks.push(embeddedChunk) + embeddingProgress.completed++ + updateProgress?.({ + completedChunks: embeddingProgress.completed, + totalChunks: contentChunks.length, + totalFiles: filesToIndex.length, + }) + + // Insert vectors in batches + if ( + embeddingChunks.length >= + embeddingProgress.inserted + batchSize || + embeddingChunks.length === contentChunks.length + ) { + await this.repository.insertVectors( + embeddingChunks.slice( + embeddingProgress.inserted, + embeddingProgress.inserted + batchSize, + ), + embeddingModel, + ) + embeddingProgress.inserted += batchSize + } }, - }, - ) + { + numOfAttempts: 5, + startingDelay: 1000, + timeMultiple: 1.5, + jitter: 'full', + retry: (error) => { + console.error(error) + const isRateLimitError = + error.status === 429 && + error.message.toLowerCase().includes('rate limit') + return !!isRateLimitError // retry only for rate limit errors + }, + }, + ) + } catch (error) { + abortController.abort() + throw error + } }), ) try { await Promise.all(tasks) } catch (error) { - console.error('Error embedding chunks:', error) - throw error + if ( + error instanceof LLMAPIKeyNotSetException || + error instanceof LLMAPIKeyInvalidException || + error instanceof LLMBaseUrlNotSetException + ) { + openSettingsModalWithError(this.app, (error as Error).message) + } else { + console.error('Error embedding chunks:', error) + throw error + } } finally { await this.repository.save() }