diff --git a/src/LLMProviders/chainManager.ts b/src/LLMProviders/chainManager.ts new file mode 100644 index 00000000..1a482c6a --- /dev/null +++ b/src/LLMProviders/chainManager.ts @@ -0,0 +1,393 @@ +import { LangChainParams, SetChainOptions } from '@/aiParams'; +import ChainFactory, { + ChainType +} from '@/chainFactory'; +import { + AI_SENDER, + ChatModelDisplayNames +} from '@/constants'; +import { ProxyChatOpenAI } from '@/langchainWrappers'; +import { ChatMessage } from '@/sharedState'; +import { getModelName, isSupportedChain } from '@/utils'; +import VectorDBManager, { MemoryVector } from '@/vectorDBManager'; +import { ChatOllama } from "@langchain/community/chat_models/ollama"; +import { + BaseChain, + ConversationChain, + ConversationalRetrievalQAChain, + RetrievalQAChain +} from "langchain/chains"; +import { + ChatPromptTemplate, + HumanMessagePromptTemplate, + MessagesPlaceholder +} from "langchain/prompts"; +import { ContextualCompressionRetriever } from "langchain/retrievers/contextual_compression"; +import { LLMChainExtractor } from "langchain/retrievers/document_compressors/chain_extract"; +import { RecursiveCharacterTextSplitter } from "langchain/text_splitter"; +import { MemoryVectorStore } from "langchain/vectorstores/memory"; +import { Notice } from 'obsidian'; +import ChatModelManager from './chatModelManager'; +import EmbeddingsManager from './embeddingManager'; +import MemoryManager from './memoryManager'; +import PromptManager from './promptManager'; + +export default class ChainManager { + private static chain: BaseChain; + private static retrievalChain: RetrievalQAChain; + private static conversationalRetrievalChain: ConversationalRetrievalQAChain; + + private static isOllamaModelActive = false; + private static isOpenRouterModelActive = false; + + private vectorStore: MemoryVectorStore; + private promptManager: PromptManager; + private embeddingsManager: EmbeddingsManager; + public chatModelManager: ChatModelManager; + public langChainParams: LangChainParams; + public memoryManager: MemoryManager; + + /** + * Constructor for initializing langChainParams and instantiating singletons. + * + * @param {LangChainParams} langChainParams - the parameters for language chaining + * @return {void} + */ + constructor( + langChainParams: LangChainParams + ) { + // Instantiate singletons + this.langChainParams = langChainParams; + this.memoryManager = MemoryManager.getInstance(this.langChainParams); + this.chatModelManager = ChatModelManager.getInstance(this.langChainParams); + this.promptManager = PromptManager.getInstance(this.langChainParams); + this.embeddingsManager = EmbeddingsManager.getInstance(this.langChainParams); + this.createChainWithNewModel(this.langChainParams.modelDisplayName); + } + + private setNoteContent(noteContent: string): void { + this.langChainParams.options.noteContent = noteContent; + } + + private validateChainType(chainType: ChainType): void { + if (chainType === undefined || chainType === null) throw new Error('No chain type set'); + } + + /** + * Update the active model and create a new chain + * with the specified model display name. + * + * @param {string} newModelDisplayName - the display name of the new model in the dropdown + * @return {void} + */ + createChainWithNewModel(newModelDisplayName: string): void { + ChainManager.isOllamaModelActive = newModelDisplayName === ChatModelDisplayNames.OLLAMA; + ChainManager.isOpenRouterModelActive = newModelDisplayName === ChatModelDisplayNames.OPENROUTERAI; + // model and model display name must be update at the same time! + let newModel = getModelName(newModelDisplayName); + + switch (newModelDisplayName) { + case ChatModelDisplayNames.OLLAMA: + newModel = this.langChainParams.ollamaModel; + break; + case ChatModelDisplayNames.LM_STUDIO: + newModel = 'check_model_in_lm_studio_ui'; + break; + case ChatModelDisplayNames.OPENROUTERAI: + newModel = this.langChainParams.openRouterModel; + break; + } + + try { + this.langChainParams.model = newModel; + this.langChainParams.modelDisplayName = newModelDisplayName; + this.chatModelManager.setChatModel(newModelDisplayName); + // Must update the chatModel for chain because ChainFactory always + // retrieves the old chain without the chatModel change if it exists! + // Create a new chain with the new chatModel + this.createChain( + this.langChainParams.chainType, + {...this.langChainParams.options, forceNewCreation: true}, + ) + console.log(`Setting model to ${newModelDisplayName}: ${newModel}`); + } catch (error) { + console.error("createChainWithNewModel failed: ", error); + console.log("model:", this.langChainParams.model); + } + } + + /* Create a new chain, or update chain with new model */ + createChain( + chainType: ChainType, + options?: SetChainOptions, + ): void { + this.validateChainType(chainType); + try { + this.setChain(chainType, options); + } catch (error) { + new Notice('Error creating chain:', error); + console.error('Error creating chain:', error); + } + } + + async setChain( + chainType: ChainType, + options: SetChainOptions = {}, + ): Promise { + if (!this.chatModelManager.validateChatModel(this.chatModelManager.getChatModel())) { + // No need to throw error and trigger multiple Notices to user + console.error('setChain failed: No chat model set.'); + return; + } + this.validateChainType(chainType); + + // Get chatModel, memory, prompt, and embeddingAPI from respective managers + const chatModel = this.chatModelManager.getChatModel(); + const memory = this.memoryManager.getMemory(); + const chatPrompt = this.promptManager.getChatPrompt(); + const embeddingsAPI = this.embeddingsManager.getEmbeddingsAPI(); + + switch (chainType) { + case ChainType.LLM_CHAIN: { + // For initial load of the plugin + if (options.forceNewCreation) { + // setChain is async, this is to ensure Ollama has the right model passed in from the setting + if (ChainManager.isOllamaModelActive) { + (chatModel as ChatOllama).model = this.langChainParams.ollamaModel; + } else if (ChainManager.isOpenRouterModelActive) { + (chatModel as ProxyChatOpenAI).modelName = this.langChainParams.openRouterModel; + } + + ChainManager.chain = ChainFactory.createNewLLMChain({ + llm: chatModel, + memory: memory, + prompt: options.prompt || chatPrompt, + }) as ConversationChain; + } else { + // For navigating back to the plugin view + ChainManager.chain = ChainFactory.getLLMChainFromMap({ + llm: chatModel, + memory: memory, + prompt: options.prompt || chatPrompt, + }) as ConversationChain; + } + + this.langChainParams.chainType = ChainType.LLM_CHAIN; + break; + } + case ChainType.RETRIEVAL_QA_CHAIN: { + if (!options.noteContent) { + new Notice('No note content provided'); + throw new Error('No note content provided'); + } + + this.setNoteContent(options.noteContent); + const docHash = VectorDBManager.getDocumentHash(options.noteContent); + const parsedMemoryVectors: MemoryVector[] | undefined = await VectorDBManager.getMemoryVectors(docHash); + if (parsedMemoryVectors) { + const vectorStore = await VectorDBManager.rebuildMemoryVectorStore( + parsedMemoryVectors, embeddingsAPI + ); + ChainManager.retrievalChain = RetrievalQAChain.fromLLM( + chatModel, + vectorStore.asRetriever(), + ); + console.log('Existing vector store for document hash: ', docHash); + } else { + await this.buildIndex(options.noteContent, docHash); + if (!this.vectorStore) { + console.error('Error creating vector store.'); + return; + } + + const baseCompressor = LLMChainExtractor.fromLLM(chatModel); + const retriever = new ContextualCompressionRetriever({ + baseCompressor, + baseRetriever: this.vectorStore.asRetriever(), + }); + ChainManager.retrievalChain = RetrievalQAChain.fromLLM( + chatModel, + retriever, + ); + console.log( + 'New retrieval qa chain with contextual compression created for ' + + 'document hash: ', docHash + ); + } + + this.langChainParams.chainType = ChainType.RETRIEVAL_QA_CHAIN; + console.log('Set chain:', ChainType.RETRIEVAL_QA_CHAIN); + break; + } + default: + this.validateChainType(chainType); + break; + } + } + + async buildIndex(noteContent: string, docHash: string): Promise { + const textSplitter = new RecursiveCharacterTextSplitter({ chunkSize: 1000 }); + + const docs = await textSplitter.createDocuments([noteContent]); + const embeddingsAPI = this.embeddingsManager.getEmbeddingsAPI(); + + // Note: HF can give 503 errors frequently (it's free) + console.log('Creating vector store...'); + try { + this.vectorStore = await MemoryVectorStore.fromDocuments( + docs, embeddingsAPI, + ); + // Serialize and save vector store to PouchDB + VectorDBManager.setMemoryVectors(this.vectorStore.memoryVectors, docHash); + console.log('Vector store created successfully.'); + new Notice('Vector store created successfully.'); + } catch (error) { + new Notice('Failed to create vector store, please try again:', error); + console.error('Failed to create vector store, please try again.:', error); + } + } + + async runChain( + userMessage: string, + abortController: AbortController, + updateCurrentAiMessage: (message: string) => void, + addMessage: (message: ChatMessage) => void, + options: { debug?: boolean, ignoreSystemMessage?: boolean } = {}, + ) { + const { debug = false, ignoreSystemMessage = false } = options; + + // Check if chat model is initialized + if (!this.chatModelManager.validateChatModel(this.chatModelManager.getChatModel())) { + const errorMsg = 'Chat model is not initialized properly, check your API key in Copilot setting and make sure you have API access.'; + new Notice(errorMsg); + console.error(errorMsg); + return; + } + // Check if chain is initialized properly + if (!ChainManager.chain || !isSupportedChain(ChainManager.chain)) { + console.error( + 'Chain is not initialized properly, re-initializing chain: ', + this.langChainParams.chainType + ); + this.setChain(this.langChainParams.chainType, this.langChainParams.options); + } + + const { + temperature, + maxTokens, + systemMessage, + chatContextTurns, + chainType, + } = this.langChainParams; + + const memory = this.memoryManager.getMemory(); + const chatPrompt = this.promptManager.getChatPrompt(); + const systemPrompt = ignoreSystemMessage ? '' : systemMessage; + // Whether to ignore system prompt (for commands) + if (ignoreSystemMessage) { + const effectivePrompt = ignoreSystemMessage + ? ChatPromptTemplate.fromMessages([ + new MessagesPlaceholder("history"), + HumanMessagePromptTemplate.fromTemplate("{input}"), + ]) + : chatPrompt; + + this.setChain(chainType, { + ...this.langChainParams.options, + prompt: effectivePrompt, + }); + } else { + this.setChain(chainType, this.langChainParams.options); + } + + let fullAIResponse = ''; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const chain = ChainManager.chain as any; + try { + switch(chainType) { + case ChainType.LLM_CHAIN: + if (debug) { + console.log(`*** DEBUG INFO ***\n` + + `user message: ${userMessage}\n` + // ChatOpenAI has modelName, some other ChatModels like ChatOllama have model + + `model: ${chain.llm.modelName || chain.llm.model}\n` + + `chain type: ${chainType}\n` + + `temperature: ${temperature}\n` + + `maxTokens: ${maxTokens}\n` + + `system prompt: ${systemPrompt}\n` + + `chat context turns: ${chatContextTurns}\n`, + ); + console.log('chain:', chain); + console.log('Chat memory:', memory); + } + await ChainManager.chain.call( + { + input: userMessage, + signal: abortController.signal, + }, + [ + { + handleLLMNewToken: (token) => { + fullAIResponse += token; + updateCurrentAiMessage(fullAIResponse); + } + } + ] + ); + break; + case ChainType.RETRIEVAL_QA_CHAIN: + if (debug) { + console.log(`*** DEBUG INFO ***\n` + + `user message: ${userMessage}\n` + + `model: ${chain.llm.modelName}\n` + + `chain type: ${chainType}\n` + + `temperature: ${temperature}\n` + + `maxTokens: ${maxTokens}\n` + + `system prompt: ${systemPrompt}\n` + + `chat context turns: ${chatContextTurns}\n`, + ); + console.log('chain:', chain); + console.log('embedding provider:', this.langChainParams.embeddingProvider); + } + await ChainManager.retrievalChain.call( + { + query: userMessage, + signal: abortController.signal, + }, + [ + { + handleLLMNewToken: (token) => { + fullAIResponse += token; + updateCurrentAiMessage(fullAIResponse); + } + } + ] + ); + break; + default: + console.error('Chain type not supported:', this.langChainParams.chainType); + } + } catch (error) { + const errorData = error?.response?.data?.error || error; + const errorCode = errorData?.code || error; + if (errorCode === 'model_not_found') { + const modelNotFoundMsg = "You do not have access to this model or the model does not exist, please check with your API provider."; + new Notice(modelNotFoundMsg); + console.error(modelNotFoundMsg); + } else { + new Notice(`LangChain error: ${errorCode}`); + console.error(errorData); + } + } finally { + if (fullAIResponse) { + addMessage({ + message: fullAIResponse, + sender: AI_SENDER, + isVisible: true, + }); + } + updateCurrentAiMessage(''); + } + return fullAIResponse; + } +} diff --git a/src/LLMProviders/chatModelManager.ts b/src/LLMProviders/chatModelManager.ts new file mode 100644 index 00000000..ddcd4b68 --- /dev/null +++ b/src/LLMProviders/chatModelManager.ts @@ -0,0 +1,194 @@ +import { LangChainParams, ModelConfig } from '@/aiParams'; +import { + AZURE_MODELS, + GOOGLE_MODELS, + LM_STUDIO_MODELS, + ModelProviders, + OLLAMA_MODELS, + OPENAI_MODELS, + OPENROUTERAI_MODELS +} from '@/constants'; +import { ProxyChatOpenAI } from '@/langchainWrappers'; +import { ChatOllama } from "@langchain/community/chat_models/ollama"; +import { ChatGoogleGenerativeAI } from "@langchain/google-genai"; +import { BaseChatModel } from 'langchain/chat_models/base'; +import { ChatOpenAI } from 'langchain/chat_models/openai'; +import { Notice } from 'obsidian'; + +export default class ChatModelManager { + private static instance: ChatModelManager; + private static chatModel: BaseChatModel; + private static chatOpenAI: ChatOpenAI; + private static modelMap: Record< + string, + { + hasApiKey: boolean; + AIConstructor: new (config: ModelConfig) => BaseChatModel; + vendor: string; + } + >; + + + private constructor( + private langChainParams: LangChainParams + ) { + this.buildModelMap(); + } + + static getInstance( + langChainParams: LangChainParams + ): ChatModelManager { + if (!ChatModelManager.instance) { + ChatModelManager.instance = new ChatModelManager(langChainParams); + } + return ChatModelManager.instance; + } + + private getModelConfig(chatModelProvider: string): ModelConfig { + const params = this.langChainParams; + const baseConfig: ModelConfig = { + modelName: params.model, + temperature: params.temperature, + streaming: true, + maxRetries: 3, + maxConcurrency: 3, + }; + + const providerConfig = { + [ModelProviders.OPENAI]: { + openAIApiKey: params.openAIApiKey, + maxTokens: params.maxTokens, + openAIProxyBaseUrl: params.openAIProxyBaseUrl, + }, + [ModelProviders.ANTHROPIC]: { + anthropicApiKey: params.anthropicApiKey, + }, + [ModelProviders.AZURE_OPENAI]: { + maxTokens: params.maxTokens, + azureOpenAIApiKey: params.azureOpenAIApiKey, + azureOpenAIApiInstanceName: params.azureOpenAIApiInstanceName, + azureOpenAIApiDeploymentName: params.azureOpenAIApiDeploymentName, + azureOpenAIApiVersion: params.azureOpenAIApiVersion, + }, + [ModelProviders.GOOGLE]: { + apiKey: params.googleApiKey, + }, + [ModelProviders.OPENROUTERAI]: { + modelName: params.openRouterModel, + openAIApiKey: params.openRouterAiApiKey, + openAIProxyBaseUrl: 'https://openrouter.ai/api/v1', + }, + [ModelProviders.LM_STUDIO]: { + openAIApiKey: 'placeholder', + openAIProxyBaseUrl: `http://localhost:${params.lmStudioPort}/v1`, + }, + [ModelProviders.OLLAMA]: { + ...(params.ollamaBaseUrl ? { baseUrl: params.ollamaBaseUrl } : {}), + modelName: params.ollamaModel, + }, + }; + + return { ...baseConfig, ...(providerConfig[chatModelProvider as keyof typeof providerConfig] || {}) }; + } + + private buildModelMap() { + ChatModelManager.modelMap = {}; + const modelMap = ChatModelManager.modelMap; + + const OpenAIChatModel = this.langChainParams.openAIProxyBaseUrl + ? ProxyChatOpenAI : ChatOpenAI; + + const modelConfigurations = [ + { + models: OPENAI_MODELS, + apiKey: this.langChainParams.openAIApiKey, + constructor: OpenAIChatModel, + vendor: ModelProviders.OPENAI, + }, + { + models: AZURE_MODELS, + apiKey: this.langChainParams.azureOpenAIApiKey, + constructor: ChatOpenAI, + vendor: ModelProviders.AZURE_OPENAI, + }, + { + models: GOOGLE_MODELS, + apiKey: this.langChainParams.googleApiKey, + constructor: ChatGoogleGenerativeAI, + vendor: ModelProviders.GOOGLE, + }, + { + models: OPENROUTERAI_MODELS, + apiKey: this.langChainParams.openRouterAiApiKey, + constructor: ProxyChatOpenAI, + vendor: ModelProviders.OPENROUTERAI, + }, + { + models: OLLAMA_MODELS, + apiKey: true, + constructor: ChatOllama, + vendor: ModelProviders.OLLAMA, + }, + { + models: LM_STUDIO_MODELS, + apiKey: true, + constructor: ProxyChatOpenAI, + vendor: ModelProviders.LM_STUDIO, + }, + ]; + + modelConfigurations.forEach(({ models, apiKey, constructor, vendor }) => { + models.forEach(modelDisplayNameKey => { + modelMap[modelDisplayNameKey] = { + hasApiKey: Boolean(apiKey), + AIConstructor: constructor, + vendor: vendor, + }; + }); + }); + } + + getChatModel(): BaseChatModel { + return ChatModelManager.chatModel; + } + + setChatModel(modelDisplayName: string): void { + if (!ChatModelManager.modelMap.hasOwnProperty(modelDisplayName)) { + throw new Error(`No model found for: ${modelDisplayName}`); + } + + // Create and return the appropriate model + const selectedModel = ChatModelManager.modelMap[modelDisplayName]; + if (!selectedModel.hasApiKey) { + const errorMessage = `API key is not provided for the model: ${modelDisplayName}. Model switch failed.`; + new Notice(errorMessage); + // Stop execution and deliberate fail the model switch + throw new Error(errorMessage); + } + + const modelConfig = this.getModelConfig(selectedModel.vendor); + + try { + const newModelInstance = new selectedModel.AIConstructor({ + ...modelConfig, + }); + + // Set the new model + ChatModelManager.chatModel = newModelInstance; + } catch (error) { + console.error(error); + new Notice(`Error creating model: ${modelDisplayName}`); + } + } + + validateChatModel(chatModel: BaseChatModel): boolean { + if (chatModel === undefined || chatModel === null) { + return false; + } + return true + } + + async countTokens(inputStr: string): Promise { + return ChatModelManager.chatOpenAI.getNumTokens(inputStr); + } +} \ No newline at end of file diff --git a/src/LLMProviders/embeddingManager.ts b/src/LLMProviders/embeddingManager.ts new file mode 100644 index 00000000..34e12cf8 --- /dev/null +++ b/src/LLMProviders/embeddingManager.ts @@ -0,0 +1,81 @@ +import { LangChainParams } from '@/aiParams'; +import { ModelProviders } from '@/constants'; +import { ProxyOpenAIEmbeddings } from '@/langchainWrappers'; +import { CohereEmbeddings } from "@langchain/cohere"; +import { Embeddings } from "langchain/embeddings/base"; +import { HuggingFaceInferenceEmbeddings } from "langchain/embeddings/hf"; +import { OpenAIEmbeddings } from "langchain/embeddings/openai"; + +export default class EmbeddingManager { + private static instance: EmbeddingManager; + private constructor( + private langChainParams: LangChainParams + ) {} + + static getInstance( + langChainParams: LangChainParams + ): EmbeddingManager { + if (!EmbeddingManager.instance) { + EmbeddingManager.instance = new EmbeddingManager(langChainParams); + } + return EmbeddingManager.instance; + } + + getEmbeddingsAPI(): Embeddings { + const { + openAIApiKey, + azureOpenAIApiKey, + azureOpenAIApiInstanceName, + azureOpenAIApiVersion, + azureOpenAIApiEmbeddingDeploymentName, + openAIProxyBaseUrl, + } = this.langChainParams; + + // Note that openAIProxyBaseUrl has the highest priority. + // If openAIProxyBaseUrl is set, it overrides both chat and embedding models. + const OpenAIEmbeddingsAPI = openAIProxyBaseUrl ? + new ProxyOpenAIEmbeddings({ + openAIApiKey, + maxRetries: 3, + maxConcurrency: 3, + timeout: 10000, + openAIProxyBaseUrl, + }): + new OpenAIEmbeddings({ + openAIApiKey, + maxRetries: 3, + maxConcurrency: 3, + timeout: 10000, + }); + + switch(this.langChainParams.embeddingProvider) { + case ModelProviders.OPENAI: + return OpenAIEmbeddingsAPI + case ModelProviders.HUGGINGFACE: + return new HuggingFaceInferenceEmbeddings({ + apiKey: this.langChainParams.huggingfaceApiKey, + maxRetries: 3, + maxConcurrency: 3, + }); + case ModelProviders.COHEREAI: + return new CohereEmbeddings({ + apiKey: this.langChainParams.cohereApiKey, + maxRetries: 3, + maxConcurrency: 3, + }); + case ModelProviders.AZURE_OPENAI: + return new OpenAIEmbeddings({ + azureOpenAIApiKey, + azureOpenAIApiInstanceName, + azureOpenAIApiDeploymentName: azureOpenAIApiEmbeddingDeploymentName, + azureOpenAIApiVersion, + maxRetries: 3, + maxConcurrency: 3, + }); + default: + console.error('No embedding provider set. Using OpenAI.'); + return OpenAIEmbeddingsAPI; + } + } + +} \ No newline at end of file diff --git a/src/LLMProviders/memoryManager.ts b/src/LLMProviders/memoryManager.ts new file mode 100644 index 00000000..ad47a99a --- /dev/null +++ b/src/LLMProviders/memoryManager.ts @@ -0,0 +1,41 @@ +import { LangChainParams } from '@/aiParams'; +import { BufferWindowMemory } from "langchain/memory"; + +export default class MemoryManager { + private static instance: MemoryManager; + private memory: BufferWindowMemory; + + private constructor( + private langChainParams: LangChainParams + ) { + this.initMemory(); + } + + static getInstance( + langChainParams: LangChainParams + ): MemoryManager { + if (!MemoryManager.instance) { + MemoryManager.instance = new MemoryManager(langChainParams); + } + return MemoryManager.instance; + } + + private initMemory(): void { + this.memory = new BufferWindowMemory({ + k: this.langChainParams.chatContextTurns * 2, + memoryKey: 'history', + inputKey: 'input', + returnMessages: true, + }); + } + + getMemory(): BufferWindowMemory { + return this.memory; + } + + clearChatMemory(): void { + console.log('clearing chat memory'); + this.memory.clear(); + } + +} \ No newline at end of file diff --git a/src/LLMProviders/promptManager.ts b/src/LLMProviders/promptManager.ts new file mode 100644 index 00000000..452eea88 --- /dev/null +++ b/src/LLMProviders/promptManager.ts @@ -0,0 +1,41 @@ +import { LangChainParams } from '@/aiParams'; +import { + ChatPromptTemplate, + HumanMessagePromptTemplate, + MessagesPlaceholder, + SystemMessagePromptTemplate, +} from "langchain/prompts"; + +export default class PromptManager { + private static instance: PromptManager; + private chatPrompt: ChatPromptTemplate; + + private constructor( + private langChainParams: LangChainParams + ) { + this.initChatPrompt(); + } + + static getInstance( + langChainParams: LangChainParams + ): PromptManager { + if (!PromptManager.instance) { + PromptManager.instance = new PromptManager(langChainParams); + } + return PromptManager.instance; + } + + private initChatPrompt(): void { + this.chatPrompt = ChatPromptTemplate.fromMessages([ + SystemMessagePromptTemplate.fromTemplate( + this.langChainParams.systemMessage + ), + new MessagesPlaceholder("history"), + HumanMessagePromptTemplate.fromTemplate("{input}"), + ]); + } + + getChatPrompt(): ChatPromptTemplate { + return this.chatPrompt; + } +} \ No newline at end of file diff --git a/src/aiParams.ts b/src/aiParams.ts index 14918488..2f4c7a42 100644 --- a/src/aiParams.ts +++ b/src/aiParams.ts @@ -43,7 +43,7 @@ export interface LangChainParams { systemMessage: string, chatContextTurns: number, embeddingProvider: string, - chainType: ChainType, // Default ChainType is set in main.ts getAIStateParams + chainType: ChainType, // Default ChainType is set in main.ts getChainManagerParams options: SetChainOptions, ollamaModel: string, ollamaBaseUrl: string, diff --git a/src/aiState.ts b/src/aiState.ts index 4a5e11c7..7594c36c 100644 --- a/src/aiState.ts +++ b/src/aiState.ts @@ -1,805 +1,14 @@ -import ChainFactory, { - ChainType -} from '@/chainFactory'; -import { - AI_SENDER, - ANTHROPIC, - AZURE_MODELS, - AZURE_OPENAI, - COHEREAI, - ChatModelDisplayNames, - DEFAULT_SYSTEM_PROMPT, - GOOGLE, - GOOGLE_MODELS, - HUGGINGFACE, - LM_STUDIO, - LM_STUDIO_MODELS, - OLLAMA, - OLLAMA_MODELS, - OPENAI, - OPENAI_MODELS, - OPENROUTERAI, - OPENROUTERAI_MODELS, - USER_SENDER, -} from '@/constants'; -import { ChatMessage } from '@/sharedState'; -import { getModelName, isSupportedChain } from '@/utils'; -import VectorDBManager, { MemoryVector } from '@/vectorDBManager'; -import { CohereEmbeddings } from "@langchain/cohere"; -import { ChatOllama } from "@langchain/community/chat_models/ollama"; -import { ChatGoogleGenerativeAI } from "@langchain/google-genai"; -import { - BaseChain, - ConversationChain, - ConversationalRetrievalQAChain, - RetrievalQAChain -} from "langchain/chains"; -import { ChatAnthropic } from 'langchain/chat_models/anthropic'; -import { BaseChatModel } from 'langchain/chat_models/base'; -import { ChatOpenAI } from 'langchain/chat_models/openai'; -import { Embeddings } from "langchain/embeddings/base"; -import { HuggingFaceInferenceEmbeddings } from "langchain/embeddings/hf"; -import { OpenAIEmbeddings } from "langchain/embeddings/openai"; +import ChainManager from '@/LLMProviders/chainManager'; +import { SetChainOptions } from '@/aiParams'; +import { ChainType } from '@/chainFactory'; import { BufferWindowMemory } from "langchain/memory"; -import { - ChatPromptTemplate, - HumanMessagePromptTemplate, - MessagesPlaceholder, - SystemMessagePromptTemplate, -} from "langchain/prompts"; -import { ContextualCompressionRetriever } from "langchain/retrievers/contextual_compression"; -import { LLMChainExtractor } from "langchain/retrievers/document_compressors/chain_extract"; -import { AIMessage, HumanMessage, SystemMessage } from 'langchain/schema'; -import { RecursiveCharacterTextSplitter } from "langchain/text_splitter"; -import { MemoryVectorStore } from "langchain/vectorstores/memory"; -import { Notice } from 'obsidian'; import { useState } from 'react'; -import { LangChainParams, ModelConfig, SetChainOptions } from './aiParams'; -import { ProxyChatOpenAI, ProxyOpenAIEmbeddings } from './langchainWrappers'; - -/** - * AIState manages the chat model, LangChain, and related state. - * - * initMemory() - Initializes the memory buffer for the chat history. - * - * initChatPrompt() - Initializes the chat prompt template with placeholders for chat history and user input. - * - * validateChainType() - Throws error if chain type is invalid. - * - * validateChatModel() - Returns false if chat model is empty. - * - * getModelConfig() - Gets model configuration based on vendor. - * - * buildModelMap() - Builds map of available models and required info to instantiate them. - * - * getEmbeddingsAPI() - Gets the appropriate embeddings API instance based on settings. - * - * clearChatMemory() - Clears the chat history memory. - * - * setChatModel() - Creates and sets the chat model instance based on name. - * - * setModel() - Sets new model name and display name, creates new chain. - * - * createChain() - Creates new chain or updates existing with new model. - * - * setChain() - Validates and sets the chain instance of specified type. - * - * buildIndex() - Builds vector index for note content using embeddings API. - * - * countTokens() - Counts number of tokens for given input string. - * - * runChatModel() - Runs chat model and handles streaming output. - * - * runChain() - Runs the LangChain and handles streaming output. - */ -class AIState { - private static chatModel: BaseChatModel; - private static chatOpenAI: ChatOpenAI; - private static chatAnthropic: ChatAnthropic; - private static azureChatOpenAI: ChatOpenAI; - private static chatGoogleGenerativeAI: ChatGoogleGenerativeAI; - private static chatOllama: ChatOllama; - private static chain: BaseChain; - private static retrievalChain: RetrievalQAChain; - private static conversationalRetrievalChain: ConversationalRetrievalQAChain; - - private chatPrompt: ChatPromptTemplate; - private vectorStore: MemoryVectorStore; - - private static isOllamaModelActive = false; - private static isOpenRouterModelActive = false; - - memory: BufferWindowMemory; - langChainParams: LangChainParams; - modelMap: Record< - string, - { - hasApiKey: boolean; - AIConstructor: new (config: ModelConfig) => BaseChatModel; - vendor: string; - }>; - - constructor(langChainParams: LangChainParams) { - this.langChainParams = langChainParams; - this.buildModelMap(); - this.initMemory(); - this.initChatPrompt(); - // This takes care of chain creation as well - this.setModel(this.langChainParams.modelDisplayName); - } - - private initMemory(): void { - this.memory = new BufferWindowMemory({ - k: this.langChainParams.chatContextTurns * 2, - memoryKey: 'history', - inputKey: 'input', - returnMessages: true, - }); - } - - private setNoteContent(noteContent: string): void { - this.langChainParams.options.noteContent = noteContent; - } - - private initChatPrompt(): void { - this.chatPrompt = ChatPromptTemplate.fromMessages([ - SystemMessagePromptTemplate.fromTemplate( - this.langChainParams.systemMessage - ), - new MessagesPlaceholder("history"), - HumanMessagePromptTemplate.fromTemplate("{input}"), - ]); - } - - private validateChainType(chainType: ChainType): void { - if (chainType === undefined || chainType === null) throw new Error('No chain type set'); - } - - private validateChatModel(chatModel: BaseChatModel): boolean { - if (chatModel === undefined || chatModel === null) { - return false; - } - return true - } - - private getModelConfig(chatModelProvider: string): ModelConfig { - const { - openAIApiKey, - anthropicApiKey, - azureOpenAIApiKey, - azureOpenAIApiInstanceName, - azureOpenAIApiDeploymentName, - azureOpenAIApiVersion, - model, - temperature, - maxTokens, - openAIProxyBaseUrl, - googleApiKey, - openRouterAiApiKey, - ollamaModel, - ollamaBaseUrl, - openRouterModel, - } = this.langChainParams; - - // Create a base configuration that applies to all models - let config: ModelConfig = { - modelName: model, - temperature: temperature, - streaming: true, - maxRetries: 3, - maxConcurrency: 3, - }; - - switch(chatModelProvider) { - case OPENAI: - config = { - ...config, - openAIApiKey, - maxTokens, - openAIProxyBaseUrl, - }; - break; - case ANTHROPIC: - config = { - ...config, - anthropicApiKey, - }; - break; - case AZURE_OPENAI: - config = { - ...config, - maxTokens, - azureOpenAIApiKey: azureOpenAIApiKey, - azureOpenAIApiInstanceName: azureOpenAIApiInstanceName, - azureOpenAIApiDeploymentName: azureOpenAIApiDeploymentName, - azureOpenAIApiVersion: azureOpenAIApiVersion, - }; - break; - case GOOGLE: - config = { - ...config, - apiKey: googleApiKey, - }; - break; - case OPENROUTERAI: - config = { - ...config, - modelName: openRouterModel, - openAIApiKey: openRouterAiApiKey, - openAIProxyBaseUrl: 'https://openrouter.ai/api/v1', - }; - break; - case LM_STUDIO: - config = { - ...config, - openAIApiKey: 'placeholder', - openAIProxyBaseUrl: `http://localhost:${this.langChainParams.lmStudioPort}/v1`, - }; - break; - case OLLAMA: - config = { - ...config, - ...(ollamaBaseUrl ? { baseUrl: ollamaBaseUrl } : {}), - modelName: ollamaModel, - }; - break; - } - - return config; - } - - private buildModelMap() { - const modelMap: Record< - string, - { - hasApiKey: boolean; - AIConstructor: new (config: ModelConfig) => BaseChatModel; - vendor: string; - } - > = {}; - - const OpenAIChatModel = this.langChainParams.openAIProxyBaseUrl - ? ProxyChatOpenAI : ChatOpenAI; - - // Build modelMap - for (const modelDisplayNameKey of OPENAI_MODELS) { - modelMap[modelDisplayNameKey] = { - hasApiKey: Boolean(this.langChainParams.openAIApiKey), - AIConstructor: OpenAIChatModel, - vendor: OPENAI, - }; - } - - for (const modelDisplayNameKey of AZURE_MODELS) { - modelMap[modelDisplayNameKey] = { - hasApiKey: Boolean(this.langChainParams.azureOpenAIApiKey), - AIConstructor: ChatOpenAI, - vendor: AZURE_OPENAI, - }; - } - - for (const modelDisplayNameKey of GOOGLE_MODELS) { - modelMap[modelDisplayNameKey] = { - hasApiKey: Boolean(this.langChainParams.googleApiKey), - AIConstructor: ChatGoogleGenerativeAI, - vendor: GOOGLE, - }; - } - - for (const modelDisplayNameKey of OPENROUTERAI_MODELS) { - modelMap[modelDisplayNameKey] = { - hasApiKey: Boolean(this.langChainParams.openRouterAiApiKey), - AIConstructor: ProxyChatOpenAI, - vendor: OPENROUTERAI, - }; - } - - for (const modelDisplayNameKey of OLLAMA_MODELS) { - modelMap[modelDisplayNameKey] = { - hasApiKey: true, - AIConstructor: ChatOllama, - vendor: OLLAMA, - }; - } - - for (const modelDisplayNameKey of LM_STUDIO_MODELS) { - modelMap[modelDisplayNameKey] = { - hasApiKey: true, - AIConstructor: ProxyChatOpenAI, - vendor: LM_STUDIO, - }; - } - - this.modelMap = modelMap; - } - - getEmbeddingsAPI(): Embeddings { - const { - openAIApiKey, - azureOpenAIApiKey, - azureOpenAIApiInstanceName, - azureOpenAIApiVersion, - azureOpenAIApiEmbeddingDeploymentName, - openAIProxyBaseUrl, - } = this.langChainParams; - - // Note that openAIProxyBaseUrl has the highest priority. - // If openAIProxyBaseUrl is set, it overrides both chat and embedding models. - const OpenAIEmbeddingsAPI = openAIProxyBaseUrl ? - new ProxyOpenAIEmbeddings({ - openAIApiKey, - maxRetries: 3, - maxConcurrency: 3, - timeout: 10000, - openAIProxyBaseUrl, - }): - new OpenAIEmbeddings({ - openAIApiKey, - maxRetries: 3, - maxConcurrency: 3, - timeout: 10000, - }); - - switch(this.langChainParams.embeddingProvider) { - case OPENAI: - // Every OpenAIEmbedding call is giving a 'refused to set header user-agent' - // It's generally not a problem. - // TODO: Check if this error can be avoided. - return OpenAIEmbeddingsAPI - case HUGGINGFACE: - // TODO: This does not have a timeout param, need to check in the future. - return new HuggingFaceInferenceEmbeddings({ - apiKey: this.langChainParams.huggingfaceApiKey, - maxRetries: 3, - maxConcurrency: 3, - }); - case COHEREAI: - return new CohereEmbeddings({ - apiKey: this.langChainParams.cohereApiKey, - maxRetries: 3, - maxConcurrency: 3, - }); - case AZURE_OPENAI: - return new OpenAIEmbeddings({ - azureOpenAIApiKey, - azureOpenAIApiInstanceName, - azureOpenAIApiDeploymentName: azureOpenAIApiEmbeddingDeploymentName, - azureOpenAIApiVersion, - maxRetries: 3, - maxConcurrency: 3, - }); - // TODO: Check Ollama local embedding and come back here - // case OLLAMA: - // return new ProxyOpenAIEmbeddings({ - // openAIApiKey, - // openAIProxyBaseUrl, - // maxRetries: 3, - // maxConcurrency: 3, - // timeout: 10000, - // }) - default: - console.error('No embedding provider set. Using OpenAI.'); - return OpenAIEmbeddingsAPI; - } - } - - clearChatMemory(): void { - console.log('clearing chat memory'); - this.memory.clear(); - } - - private setChatModel(modelDisplayName: string): void { - if (!this.modelMap.hasOwnProperty(modelDisplayName)) { - throw new Error(`No model found for: ${modelDisplayName}`); - } - - // Create and return the appropriate model - const selectedModel = this.modelMap[modelDisplayName]; - if (!selectedModel.hasApiKey) { - const errorMessage = `API key is not provided for the model: ${modelDisplayName}. Model switch failed.`; - new Notice(errorMessage); - // Stop execution and deliberate fail the model switch - throw new Error(errorMessage); - } - - const modelConfig = this.getModelConfig(selectedModel.vendor); - - try { - const newModelInstance = new selectedModel.AIConstructor({ - ...modelConfig, - }); - - switch(selectedModel.vendor) { - case OPENAI: - AIState.chatOpenAI = newModelInstance as ChatOpenAI; - break; - case ANTHROPIC: - AIState.chatAnthropic = newModelInstance as ChatAnthropic; - break; - case AZURE_OPENAI: - AIState.azureChatOpenAI = newModelInstance as ChatOpenAI; - break; - case GOOGLE: - AIState.chatGoogleGenerativeAI = newModelInstance as ChatGoogleGenerativeAI; - break; - case OPENROUTERAI: - AIState.chatOpenAI = newModelInstance as ProxyChatOpenAI; - break; - case OLLAMA: - AIState.chatOllama = newModelInstance as ChatOllama; - break; - } - - AIState.chatModel = newModelInstance; - } catch (error) { - console.error(error); - new Notice(`Error creating model: ${modelDisplayName}`); - } - } - - setModel(newModelDisplayName: string): void { - AIState.isOllamaModelActive = newModelDisplayName === ChatModelDisplayNames.OLLAMA; - AIState.isOpenRouterModelActive = newModelDisplayName === ChatModelDisplayNames.OPENROUTERAI; - // model and model display name must be update at the same time! - let newModel = getModelName(newModelDisplayName); - - switch (newModelDisplayName) { - case ChatModelDisplayNames.OLLAMA: - newModel = this.langChainParams.ollamaModel; - break; - case ChatModelDisplayNames.LM_STUDIO: - newModel = 'check_model_in_lm_studio_ui'; - break; - case ChatModelDisplayNames.OPENROUTERAI: - newModel = this.langChainParams.openRouterModel; - break; - } - - try { - this.langChainParams.model = newModel; - this.langChainParams.modelDisplayName = newModelDisplayName; - this.setChatModel(newModelDisplayName); - // Must update the chatModel for chain because ChainFactory always - // retrieves the old chain without the chatModel change if it exists! - // Create a new chain with the new chatModel - this.createChain( - this.langChainParams.chainType, - {...this.langChainParams.options, forceNewCreation: true}, - ) - console.log(`Setting model to ${newModelDisplayName}: ${newModel}`); - } catch (error) { - console.error("setModel failed: ", error); - console.log("model:", this.langChainParams.model); - } - } - - /* Create a new chain, or update chain with new model */ - createChain( - chainType: ChainType, - options?: SetChainOptions, - ): void { - this.validateChainType(chainType); - try { - this.setChain(chainType, options); - } catch (error) { - new Notice('Error creating chain:', error); - console.error('Error creating chain:', error); - } - } - - async setChain( - chainType: ChainType, - options: SetChainOptions = {}, - ): Promise { - if (!this.validateChatModel(AIState.chatModel)) { - // No need to throw error and trigger multiple Notices to user - console.error('setChain failed: No chat model set.'); - return; - } - this.validateChainType(chainType); - - switch (chainType) { - case ChainType.LLM_CHAIN: { - // For initial load of the plugin - if (options.forceNewCreation) { - // setChain is async, this is to ensure Ollama has the right model passed in from the setting - if (AIState.isOllamaModelActive) { - (AIState.chatModel as ChatOllama).model = this.langChainParams.ollamaModel; - } else if (AIState.isOpenRouterModelActive) { - (AIState.chatModel as ProxyChatOpenAI).modelName = this.langChainParams.openRouterModel; - } - - AIState.chain = ChainFactory.createNewLLMChain({ - llm: AIState.chatModel, - memory: this.memory, - prompt: options.prompt || this.chatPrompt, - }) as ConversationChain; - } else { - // For navigating back to the plugin view - AIState.chain = ChainFactory.getLLMChainFromMap({ - llm: AIState.chatModel, - memory: this.memory, - prompt: options.prompt || this.chatPrompt, - }) as ConversationChain; - } - - this.langChainParams.chainType = ChainType.LLM_CHAIN; - break; - } - case ChainType.RETRIEVAL_QA_CHAIN: { - if (!options.noteContent) { - new Notice('No note content provided'); - throw new Error('No note content provided'); - } - - this.setNoteContent(options.noteContent); - const docHash = VectorDBManager.getDocumentHash(options.noteContent); - const parsedMemoryVectors: MemoryVector[] | undefined = await VectorDBManager.getMemoryVectors(docHash); - if (parsedMemoryVectors) { - const vectorStore = await VectorDBManager.rebuildMemoryVectorStore( - parsedMemoryVectors, this.getEmbeddingsAPI() - ); - AIState.retrievalChain = RetrievalQAChain.fromLLM( - AIState.chatModel, - vectorStore.asRetriever(), - ); - console.log('Existing vector store for document hash: ', docHash); - } else { - await this.buildIndex(options.noteContent, docHash); - if (!this.vectorStore) { - console.error('Error creating vector store.'); - return; - } - - const baseCompressor = LLMChainExtractor.fromLLM(AIState.chatModel); - const retriever = new ContextualCompressionRetriever({ - baseCompressor, - baseRetriever: this.vectorStore.asRetriever(), - }); - AIState.retrievalChain = RetrievalQAChain.fromLLM( - AIState.chatModel, - retriever, - ); - console.log( - 'New retrieval qa chain with contextual compression created for ' - + 'document hash: ', docHash - ); - } - - this.langChainParams.chainType = ChainType.RETRIEVAL_QA_CHAIN; - console.log('Set chain:', ChainType.RETRIEVAL_QA_CHAIN); - break; - } - default: - this.validateChainType(chainType); - break; - } - } - - async buildIndex(noteContent: string, docHash: string): Promise { - const textSplitter = new RecursiveCharacterTextSplitter({ chunkSize: 1000 }); - - const docs = await textSplitter.createDocuments([noteContent]); - const embeddingsAPI = this.getEmbeddingsAPI(); - - // Note: HF can give 503 errors frequently (it's free) - console.log('Creating vector store...'); - try { - this.vectorStore = await MemoryVectorStore.fromDocuments( - docs, embeddingsAPI, - ); - // Serialize and save vector store to PouchDB - VectorDBManager.setMemoryVectors(this.vectorStore.memoryVectors, docHash); - console.log('Vector store created successfully.'); - new Notice('Vector store created successfully.'); - } catch (error) { - new Notice('Failed to create vector store, please try again:', error); - console.error('Failed to create vector store, please try again.:', error); - } - } - - async countTokens(inputStr: string): Promise { - return AIState.chatOpenAI.getNumTokens(inputStr); - } - - async runChatModel( - userMessage: ChatMessage, - chatContext: ChatMessage[], - abortController: AbortController, - updateCurrentAiMessage: (message: string) => void, - addMessage: (message: ChatMessage) => void, - debug = false, - ) { - if (debug) { - console.log('chatModel:', AIState.chatModel); - for (const [i, chatMessage] of chatContext.entries()) { - console.log( - `chat message ${i}:\nsender: ${chatMessage.sender}\n${chatMessage.message}` - ); - } - } - - const systemMessage = this.langChainParams.systemMessage || DEFAULT_SYSTEM_PROMPT; - const messages = [ - new SystemMessage(systemMessage), - ...chatContext.map((chatMessage) => { - return chatMessage.sender === USER_SENDER - ? new HumanMessage(chatMessage.message) - : new AIMessage(chatMessage.message); - }), - new HumanMessage(userMessage.message), - ]; - - let fullAIResponse = ''; - await AIState.chatModel.call( - messages, - { signal: abortController.signal }, - [ - { - handleLLMNewToken: (token) => { - fullAIResponse += token; - updateCurrentAiMessage(fullAIResponse); - } - } - ] - ); - - addMessage({ - message: fullAIResponse, - sender: AI_SENDER, - isVisible: true, - }); - updateCurrentAiMessage(''); - return fullAIResponse; - } - - async runChain( - userMessage: string, - abortController: AbortController, - updateCurrentAiMessage: (message: string) => void, - addMessage: (message: ChatMessage) => void, - options: { debug?: boolean, ignoreSystemMessage?: boolean } = {}, - ) { - const { debug = false, ignoreSystemMessage = false } = options; - - // Check if chat model is initialized - if (!this.validateChatModel(AIState.chatModel)) { - const errorMsg = 'Chat model is not initialized properly, check your API key in Copilot setting and make sure you have API access.'; - new Notice(errorMsg); - console.error(errorMsg); - return; - } - // Check if chain is initialized properly - if (!AIState.chain || !isSupportedChain(AIState.chain)) { - console.error( - 'Chain is not initialized properly, re-initializing chain: ', - this.langChainParams.chainType - ); - this.setChain(this.langChainParams.chainType, this.langChainParams.options); - } - - const { - temperature, - maxTokens, - systemMessage, - chatContextTurns, - chainType, - } = this.langChainParams; - - const systemPrompt = ignoreSystemMessage ? '' : systemMessage; - // Whether to ignore system prompt (for commands) - if (ignoreSystemMessage) { - const effectivePrompt = ignoreSystemMessage - ? ChatPromptTemplate.fromMessages([ - new MessagesPlaceholder("history"), - HumanMessagePromptTemplate.fromTemplate("{input}"), - ]) - : this.chatPrompt; - - this.setChain(chainType, { - ...this.langChainParams.options, - prompt: effectivePrompt, - }); - } else { - this.setChain(chainType, this.langChainParams.options); - } - - let fullAIResponse = ''; - const chain = AIState.chain as any; - try { - switch(chainType) { - case ChainType.LLM_CHAIN: - if (debug) { - console.log(`*** DEBUG INFO ***\n` - + `user message: ${userMessage}\n` - // ChatOpenAI has modelName, some other ChatModels like ChatOllama have model - + `model: ${chain.llm.modelName || chain.llm.model}\n` - + `chain type: ${chainType}\n` - + `temperature: ${temperature}\n` - + `maxTokens: ${maxTokens}\n` - + `system prompt: ${systemPrompt}\n` - + `chat context turns: ${chatContextTurns}\n`, - ); - console.log('chain:', chain); - console.log('Chat memory:', this.memory); - } - await AIState.chain.call( - { - input: userMessage, - signal: abortController.signal, - }, - [ - { - handleLLMNewToken: (token) => { - fullAIResponse += token; - updateCurrentAiMessage(fullAIResponse); - } - } - ] - ); - break; - case ChainType.RETRIEVAL_QA_CHAIN: - if (debug) { - console.log(`*** DEBUG INFO ***\n` - + `user message: ${userMessage}\n` - + `model: ${chain.llm.modelName}\n` - + `chain type: ${chainType}\n` - + `temperature: ${temperature}\n` - + `maxTokens: ${maxTokens}\n` - + `system prompt: ${systemPrompt}\n` - + `chat context turns: ${chatContextTurns}\n`, - ); - console.log('chain:', chain); - console.log('embedding provider:', this.langChainParams.embeddingProvider); - } - await AIState.retrievalChain.call( - { - query: userMessage, - signal: abortController.signal, - }, - [ - { - handleLLMNewToken: (token) => { - fullAIResponse += token; - updateCurrentAiMessage(fullAIResponse); - } - } - ] - ); - break; - default: - console.error('Chain type not supported:', this.langChainParams.chainType); - } - } catch (error) { - const errorData = error?.response?.data?.error || error; - const errorCode = errorData?.code || error; - if (errorCode === 'model_not_found') { - const modelNotFoundMsg = "You do not have access to this model or the model does not exist, please check with your API provider."; - new Notice(modelNotFoundMsg); - console.error(modelNotFoundMsg); - } else { - new Notice(`LangChain error: ${errorCode}`); - console.error(errorData); - } - } finally { - if (fullAIResponse) { - addMessage({ - message: fullAIResponse, - sender: AI_SENDER, - isVisible: true, - }); - } - updateCurrentAiMessage(''); - } - return fullAIResponse; - } -} /** * React hook to manage state related to model, chain and memory in Chat component. */ export function useAIState( - aiState: AIState, + chainManager: ChainManager, ): [ string, (model: string) => void, @@ -807,23 +16,23 @@ export function useAIState( (chain: ChainType, options?: SetChainOptions) => void, () => void, ] { - const { langChainParams } = aiState; + const { langChainParams } = chainManager; const [currentModel, setCurrentModel] = useState(langChainParams.modelDisplayName); const [currentChain, setCurrentChain] = useState(langChainParams.chainType); - const [, setChatMemory] = useState(aiState.memory); + const [, setChatMemory] = useState(chainManager.memoryManager.getMemory()); const clearChatMemory = () => { - aiState.clearChatMemory(); - setChatMemory(aiState.memory); + chainManager.memoryManager.clearChatMemory(); + setChatMemory(chainManager.memoryManager.getMemory()); }; const setModel = (newModelDisplayName: string) => { - aiState.setModel(newModelDisplayName); + chainManager.createChainWithNewModel(newModelDisplayName); setCurrentModel(newModelDisplayName); }; const setChain = (newChain: ChainType, options?: SetChainOptions) => { - aiState.setChain(newChain, options); + chainManager.setChain(newChain, options); setCurrentChain(newChain); }; @@ -835,5 +44,3 @@ export function useAIState( clearChatMemory, ]; } - -export default AIState; diff --git a/src/components/Chat.tsx b/src/components/Chat.tsx index aef8f72e..0631680d 100644 --- a/src/components/Chat.tsx +++ b/src/components/Chat.tsx @@ -1,4 +1,5 @@ -import AIState, { useAIState } from '@/aiState'; +import ChainManager from '@/LLMProviders/chainManager'; +import { useAIState } from '@/aiState'; import { ChainType } from '@/chainFactory'; import ChatIcons from '@/components/ChatComponents/ChatIcons'; import ChatInput from '@/components/ChatComponents/ChatInput'; @@ -17,7 +18,6 @@ import { fillInSelectionForCustomPrompt, fixGrammarSpellingSelectionPrompt, formatDateTime, - getChatContext, getFileContent, getFileName, glossaryPrompt, @@ -30,7 +30,7 @@ import { sendNoteContentPrompt, simplifyPrompt, summarizePrompt, - tocPrompt, + tocPrompt } from '@/utils'; import VectorDBManager from '@/vectorDBManager'; import { EventEmitter } from 'events'; @@ -49,7 +49,7 @@ interface CreateEffectOptions { interface ChatProps { sharedState: SharedState; - aiState: AIState; + chainManager: ChainManager; emitter: EventEmitter; getChatVisibility: () => Promise; defaultSaveFolder: string; @@ -57,22 +57,20 @@ interface ChatProps { } const Chat: React.FC = ({ - sharedState, aiState, emitter, getChatVisibility, defaultSaveFolder, debug + sharedState, chainManager, emitter, getChatVisibility, defaultSaveFolder, debug }) => { const [ chatHistory, addMessage, clearMessages, ] = useSharedState(sharedState); const [ currentModel, setModel, currentChain, setChain, clearChatMemory, - ] = useAIState(aiState); + ] = useAIState(chainManager); const [currentAiMessage, setCurrentAiMessage] = useState(''); const [inputMessage, setInputMessage] = useState(''); const [abortController, setAbortController] = useState(null); const app = useContext(AppContext); - const chatContext = getChatContext(chatHistory, aiState.langChainParams.chatContextTurns * 2); - const handleSendMessage = async () => { if (!inputMessage) return; @@ -90,8 +88,7 @@ const Chat: React.FC = ({ await getAIResponse( userMessage, - chatContext, - aiState, + chainManager, addMessage, setCurrentAiMessage, setAbortController, @@ -173,8 +170,7 @@ const Chat: React.FC = ({ await getAIResponse( promptMessageHidden, - chatContext, - aiState, + chainManager, addMessage, setCurrentAiMessage, setAbortController, @@ -203,7 +199,7 @@ const Chat: React.FC = ({ } const docHash = VectorDBManager.getDocumentHash(noteContent); - await aiState.buildIndex(noteContent, docHash); + await chainManager.buildIndex(noteContent, docHash); const activeNoteOnMessage: ChatMessage = { sender: AI_SENDER, message: `Indexing [[${noteName}]]...\n\n Please switch to "QA" in Mode Selection to ask questions about it.`, @@ -231,7 +227,7 @@ const Chat: React.FC = ({ useEffect(() => { async function handleSelection(selectedText: string) { const wordCount = selectedText.split(' ').length; - const tokenCount = await aiState.countTokens(selectedText); + const tokenCount = await chainManager.chatModelManager.countTokens(selectedText); const tokenCountMessage: ChatMessage = { sender: AI_SENDER, message: `The selected text contains ${wordCount} words and ${tokenCount} tokens.`, @@ -273,15 +269,14 @@ const Chat: React.FC = ({ } // Have a hardcoded custom temperature for some commands that need more strictness - aiState.langChainParams = { - ...aiState.langChainParams, + chainManager.langChainParams = { + ...chainManager.langChainParams, ...(custom_temperature && { temperature: custom_temperature }), }; await getAIResponse( promptMessage, - [], - aiState, + chainManager, addMessage, setCurrentAiMessage, setAbortController, diff --git a/src/components/CopilotView.tsx b/src/components/CopilotView.tsx index f24113ef..a38045cc 100644 --- a/src/components/CopilotView.tsx +++ b/src/components/CopilotView.tsx @@ -1,4 +1,4 @@ -import AIState from '@/aiState'; +import ChainManager from '@/LLMProviders/chainManager'; import Chat from '@/components/Chat'; import { CHAT_VIEWTYPE } from '@/constants'; import { AppContext } from '@/context'; @@ -12,7 +12,7 @@ import { Root, createRoot } from 'react-dom/client'; export default class CopilotView extends ItemView { private sharedState: SharedState; - private aiState: AIState; + private chainManager: ChainManager; private root: Root | null = null; private defaultSaveFolder: string; private debug = false; @@ -23,7 +23,7 @@ export default class CopilotView extends ItemView { super(leaf); this.sharedState = plugin.sharedState; this.app = plugin.app; - this.aiState = plugin.aiState; + this.chainManager = plugin.chainManager; this.debug = plugin.settings.debug; this.emitter = new EventEmitter(); this.getChatVisibility = this.getChatVisibility.bind(this); @@ -63,7 +63,7 @@ export default class CopilotView extends ItemView { = { [ChatModelDisplayNames.GPT_4]: ChatModels.GPT_4, [ChatModelDisplayNames.GPT_4_TURBO]: ChatModels.GPT_4_TURBO, [ChatModelDisplayNames.GPT_4_32K]: ChatModels.GPT_4_32K, - // [ChatModelDisplayNames.CLAUDE_1]: ChatModels.CLAUDE_1, - // [ChatModelDisplayNames.CLAUDE_1_100K]: ChatModels.CLAUDE_1_100K, - // [ChatModelDisplayNames.CLAUDE_INSTANT_1]: ChatModels.CLAUDE_INSTANT_1, - // [ChatModelDisplayNames.CLAUDE_INSTANT_1_100K]: ChatModels.CLAUDE_INSTANT_1_100K, [ChatModelDisplayNames.AZURE_GPT_35_TURBO]: ChatModels.AZURE_GPT_35_TURBO, [ChatModelDisplayNames.AZURE_GPT_35_TURBO_16K]: ChatModels.AZURE_GPT_35_TURBO_16K, [ChatModelDisplayNames.AZURE_GPT_4]: ChatModels.GPT_4, @@ -98,31 +79,32 @@ export const DISPLAY_NAME_TO_MODEL: Record = { }; // Model Providers -export const OPENAI = 'openai'; -export const HUGGINGFACE = 'huggingface'; -export const COHEREAI = 'cohereai'; -export const AZURE_OPENAI = 'azure_openai'; -export const ANTHROPIC = 'anthropic'; -export const GOOGLE = 'google'; -export const OPENROUTERAI = 'openrouterai'; -export const LM_STUDIO = 'lm_studio'; -export const OLLAMA = 'ollama'; +export enum ModelProviders { + OPENAI = 'openai', + HUGGINGFACE = 'huggingface', + COHEREAI = 'cohereai', + AZURE_OPENAI = 'azure_openai', + ANTHROPIC = 'anthropic', + GOOGLE = 'google', + OPENROUTERAI = 'openrouterai', + LM_STUDIO = 'lm_studio', + OLLAMA = 'ollama', +} export const VENDOR_MODELS: Record> = { - [OPENAI]: OPENAI_MODELS, - [AZURE_OPENAI]: AZURE_MODELS, - // [ANTHROPIC]: CLAUDE_MODELS, - [GOOGLE]: GOOGLE_MODELS, - [OPENROUTERAI]: OPENROUTERAI_MODELS, - [OLLAMA]: OLLAMA_MODELS, - [LM_STUDIO]: LM_STUDIO_MODELS, + [ModelProviders.OPENAI]: OPENAI_MODELS, + [ModelProviders.AZURE_OPENAI]: AZURE_MODELS, + [ModelProviders.GOOGLE]: GOOGLE_MODELS, + [ModelProviders.OPENROUTERAI]: OPENROUTERAI_MODELS, + [ModelProviders.OLLAMA]: OLLAMA_MODELS, + [ModelProviders.LM_STUDIO]: LM_STUDIO_MODELS, }; export const EMBEDDING_PROVIDERS = [ - OPENAI, - AZURE_OPENAI, - COHEREAI, - HUGGINGFACE, + ModelProviders.OPENAI, + ModelProviders.AZURE_OPENAI, + ModelProviders.COHEREAI, + ModelProviders.HUGGINGFACE, ]; // Embedding Models @@ -155,8 +137,7 @@ export const DEFAULT_SETTINGS: CopilotSettings = { lmStudioPort: '1234', ttlDays: 30, stream: true, - embeddingProvider: OPENAI, + embeddingProvider: ModelProviders.OPENAI, defaultSaveFolder: 'copilot-conversations', debug: false, }; -export const OPEN_AI_API_URL = 'https://api.openai.com/v1/chat/completions'; diff --git a/src/langchainStream.ts b/src/langchainStream.ts index 495abd55..ed0646d4 100644 --- a/src/langchainStream.ts +++ b/src/langchainStream.ts @@ -1,4 +1,4 @@ -import AIState from '@/aiState'; +import ChainManager from '@/LLMProviders/chainManager'; import { ChatMessage } from '@/sharedState'; import { Notice } from 'obsidian'; @@ -6,8 +6,7 @@ export type Role = 'assistant' | 'user' | 'system'; export const getAIResponse = async ( userMessage: ChatMessage, - chatContext: ChatMessage[], - aiState: AIState, + chainManager: ChainManager, addMessage: (message: ChatMessage) => void, updateCurrentAiMessage: (message: string) => void, updateShouldAbort: (abortController: AbortController | null) => void, @@ -16,10 +15,7 @@ export const getAIResponse = async ( const abortController = new AbortController(); updateShouldAbort(abortController); try { - // TODO: Need to run certain models without langchain - // it will mean no retrieval qa mode for those models! - - await aiState.runChain( + await chainManager.runChain( userMessage.message, abortController, updateCurrentAiMessage, diff --git a/src/main.ts b/src/main.ts index 438c1540..bfef517e 100644 --- a/src/main.ts +++ b/src/main.ts @@ -1,5 +1,5 @@ +import ChainManager from '@/LLMProviders/chainManager'; import { LangChainParams, SetChainOptions } from '@/aiParams'; -import AIState from '@/aiState'; import { ChainType } from '@/chainFactory'; import { AddPromptModal } from "@/components/AddPromptModal"; import CopilotView from '@/components/CopilotView'; @@ -29,7 +29,7 @@ export default class CopilotPlugin extends Plugin { // A chat history that stores the messages sent and received // Only reset when the user explicitly clicks "New Chat" sharedState: SharedState; - aiState: AIState; + chainManager: ChainManager; activateViewPromise: Promise | null = null; chatIsVisible = false; dbPrompts: PouchDB.Database; @@ -41,10 +41,10 @@ export default class CopilotPlugin extends Plugin { async onload(): Promise { await this.loadSettings(); this.addSettingTab(new CopilotSettingTab(this.app, this)); - // Always have one instance of sharedState and aiState in the plugin + // Always have one instance of sharedState and chainManager in the plugin this.sharedState = new SharedState(); - const langChainParams = this.getAIStateParams(); - this.aiState = new AIState(langChainParams); + const langChainParams = this.getChainManagerParams(); + this.chainManager = new ChainManager(langChainParams); this.dbPrompts = new PouchDB('copilot_custom_prompts'); this.dbVectorStores = new PouchDB('copilot_vector_stores'); @@ -447,7 +447,7 @@ export default class CopilotPlugin extends Plugin { .filter(Boolean) as string[]; } - getAIStateParams(): LangChainParams { + getChainManagerParams(): LangChainParams { const { openAIApiKey, huggingfaceApiKey, diff --git a/src/openAiStream.ts b/src/openAiStream.ts deleted file mode 100644 index 3a2d61fd..00000000 --- a/src/openAiStream.ts +++ /dev/null @@ -1,255 +0,0 @@ -import { - AI_SENDER, - DEFAULT_SYSTEM_PROMPT, - OPEN_AI_API_URL, USER_SENDER, -} from '@/constants'; -import { ChatMessage } from '@/sharedState'; -import { Notice, requestUrl } from 'obsidian'; -import { SSE } from 'sse'; - -export type Role = 'assistant' | 'user' | 'system'; - -export interface OpenAiMessage { - role: Role; - content: string; -} - -export interface OpenAiParams { - model: string, - key: string, - temperature: number, - maxTokens: number, -} - -export class OpenAIRequestManager { - stopRequested = false; - - constructor() {} - - stopStreaming() { - this.stopRequested = true; - } - - streamSSE = async ( - openAiParams: OpenAiParams, - messages: OpenAiMessage[], - updateCurrentAiMessage: (message: string) => void, - addMessage: (message: ChatMessage) => void, - systemPrompt: string, - ) => { - return new Promise((resolve, reject) => { - try { - const { - key, - model, - temperature, - maxTokens, - } = openAiParams; - - const formattedMessages = [ - { - role: 'system', - content: systemPrompt, - }, - ...messages, - ]; - - const url = OPEN_AI_API_URL; - const options = { - model, - messages: formattedMessages, - max_tokens: maxTokens, - temperature: temperature, - stream: true, - }; - - const source = new SSE(url, { - headers: { - 'Content-Type': 'application/json', - Authorization: `Bearer ${key ? key : process.env.OPENAI_API_KEY}`, - ...(process.env.OPENAI_ORGANIZATION && { - 'OpenAI-Organization': process.env.OPENAI_ORGANIZATION, - }), - }, - method: 'POST', - payload: JSON.stringify(options), - }); - - let aiResponse = ''; - - const addAiMessageToChatHistory = (aiResponse: string) => { - const botMessage: ChatMessage = { - message: aiResponse, - sender: AI_SENDER, - isVisible: true, - }; - addMessage(botMessage); - updateCurrentAiMessage(''); - } - - const onMessage = async (e: any) => { - if (this.stopRequested) { - console.log('Manually closing SSE stream due to stop request.'); - source.close(); - addAiMessageToChatHistory(aiResponse); - this.stopRequested = false; - resolve(null); - return; - } - - if (e.data !== "[DONE]") { - const payload = JSON.parse(e.data); - const text = payload.choices[0].delta.content; - if (!text) { - return; - } - aiResponse += text; - updateCurrentAiMessage(aiResponse); - } else { - source.close(); - addAiMessageToChatHistory(aiResponse); - resolve(aiResponse); - } - }; - - source.addEventListener('message', onMessage); - - source.addEventListener('error', (e: any) => { - source.close(); - reject(e); - }); - - source.stream(); - } catch (err) { - reject(err); - } - }); - }; -} - -export const OpenAIRequest = async ( - model: string, - key: string, - messages: OpenAiMessage[], - temperature: number, - maxTokens: number, - systemPrompt: string, -): Promise => { - const formattedMessages: OpenAiMessage[] = [ - { - role: 'system', - content: systemPrompt, - }, - ...messages, - ]; - const res = await requestUrl({ - url: OPEN_AI_API_URL, - headers: { - 'Content-Type': 'application/json', - Authorization: `Bearer ${key ? key : process.env.OPENAI_API_KEY}`, - ...(process.env.OPENAI_ORGANIZATION && { - 'OpenAI-Organization': process.env.OPENAI_ORGANIZATION, - }), - }, - method: 'POST', - body: JSON.stringify({ - model, - messages: formattedMessages, - max_tokens: maxTokens, - temperature: temperature, - stream: false, - }), - }); - - if (res.status !== 200) { - throw new Error(`OpenAI API returned an error: ${res.status}`); - } - - return res.json.choices[0].message.content; -}; - -export const getAIResponse = async ( - userMessage: ChatMessage, - chatContext: ChatMessage[], - openAiParams: OpenAiParams, - streamManager: OpenAIRequestManager, - updateCurrentAiMessage: (message: string) => void, - addMessage: (message: ChatMessage) => void, - stream = true, - debug = false, - userSystemPrompt?: string, -) => { - const { - key, - model, - temperature, - maxTokens, - } = openAiParams; - - const messages: OpenAiMessage[] = [ - ...chatContext.map((chatMessage) => { - return { - role: chatMessage.sender === USER_SENDER - ? 'user' as Role : 'assistant' as Role, - content: chatMessage.message, - }; - }), - { role: 'user', content: userMessage.message }, - ]; - - const systemPrompt = userSystemPrompt || DEFAULT_SYSTEM_PROMPT; - if (debug) { - console.log('openAiParams:', openAiParams); - console.log('stream:', stream); - console.log('system prompt:', systemPrompt); - for (const [i, message] of messages.entries()) { - console.log(`Message ${i}:\nrole: ${message.role}\n${message.content}`); - } - } - - if (stream) { - // Use streamManager.streamSSE to send message to AI and get a response - try { - await streamManager.streamSSE( - openAiParams, - messages, - updateCurrentAiMessage, - addMessage, - systemPrompt, - ); - } catch (error) { - const errorData = JSON.parse(error.data); - if (errorData && errorData.error) { - new Notice( - `OpenAI error: ${errorData.error.code}. ` - + `Pls check the console for the full error message.` - ); - } - console.error('Error in streamSSE:', error.data); - } - } else { - // Non-streaming setup using OpenAIRequest - try { - const aiResponse = await OpenAIRequest( - model, - key, - messages, - temperature, - maxTokens, - systemPrompt, - ); - - const botMessage: ChatMessage = { - message: aiResponse, - sender: AI_SENDER, - isVisible: true, - }; - addMessage(botMessage); - updateCurrentAiMessage(''); - - } catch (error) { - new Notice(`OpenAI non-streaming error: ${error.status}`); - console.error(error); - } - } -}; diff --git a/src/utils.ts b/src/utils.ts index 241af269..95e8e3f6 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -1,11 +1,7 @@ import { ChainType } from '@/chainFactory'; import { - AZURE_MODELS, - AZURE_OPENAI, DEFAULT_SETTINGS, DISPLAY_NAME_TO_MODEL, - OPENAI, - OPENAI_MODELS, USER_SENDER } from '@/constants'; import { CopilotSettings } from '@/settings/SettingsPage'; @@ -45,19 +41,6 @@ export const getModelName = (modelDisplayName: string): string => { return DISPLAY_NAME_TO_MODEL[modelDisplayName]; } -export const getModelVendorMap = (): Record => { - const model_to_vendor: Record = {}; - - for (const model of OPENAI_MODELS) { - model_to_vendor[model] = OPENAI; - } - - for (const model of AZURE_MODELS) { - model_to_vendor[model] = AZURE_OPENAI; - } - return model_to_vendor; -} - // Returns the last N messages from the chat history, // last one being the newest ai message export const getChatContext = (chatHistory: ChatMessage[], contextSize: number) => {