Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Handle llm setting errors during embedding generation #51

Merged
merged 1 commit into from
Oct 27, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 7 additions & 16 deletions src/components/chat-view/Chat.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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'
Expand Down Expand Up @@ -261,14 +261,9 @@ const Chat = forwardRef<ChatRef, ChatProps>((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)
Expand Down Expand Up @@ -324,14 +319,10 @@ const Chat = forwardRef<ChatRef, ChatProps>((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)
Expand Down
24 changes: 24 additions & 0 deletions src/utils/embedding.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 = (
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down
4 changes: 2 additions & 2 deletions src/utils/llm/exception.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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'
}
}
6 changes: 3 additions & 3 deletions src/utils/llm/ollama.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
} 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 {
Expand All @@ -30,7 +30,7 @@
const headers = req.req.headers as Record<string, string>
Object.keys(headers).forEach((k) => {
if (k.startsWith('x-stainless')) {
delete headers[k]

Check warning on line 33 in src/utils/llm/ollama.ts

View workflow job for this annotation

GitHub Actions / check

Do not delete dynamically computed property keys
}
})
return req
Expand Down Expand Up @@ -59,7 +59,7 @@
options?: LLMOptions,
): Promise<LLMResponseNonStreaming> {
if (!this.ollamaBaseUrl) {
throw new LLMABaseUrlNotSetException(
throw new LLMBaseUrlNotSetException(
'Ollama Address is missing. Please set it in settings menu.',
)
}
Expand All @@ -70,7 +70,7 @@
options?: LLMOptions,
): Promise<AsyncIterable<LLMResponseStreaming>> {
if (!this.ollamaBaseUrl) {
throw new LLMABaseUrlNotSetException(
throw new LLMBaseUrlNotSetException(
'Ollama Address is missing. Please set it in settings menu.',
)
}
Expand Down
12 changes: 12 additions & 0 deletions src/utils/openSettingsModal.ts
Original file line number Diff line number Diff line change
@@ -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()
}
121 changes: 72 additions & 49 deletions src/utils/vector-db/manager.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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'

Expand Down Expand Up @@ -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()
}
Expand Down
Loading