mirror of
https://github.com/logancyang/obsidian-copilot.git
synced 2026-07-22 07:50:24 +00:00
Refactor AIState into separate managers (#250)
This commit is contained in:
Logan Yang
GitHub
parent
f84a3bc56b
commit
373c2b8f27
14 changed files with 809 additions and 1152 deletions
393
src/LLMProviders/chainManager.ts
Normal file
393
src/LLMProviders/chainManager.ts
Normal file
|
|
@ -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<void> {
|
||||
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<void> {
|
||||
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;
|
||||
}
|
||||
}
|
||||
194
src/LLMProviders/chatModelManager.ts
Normal file
194
src/LLMProviders/chatModelManager.ts
Normal file
|
|
@ -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<number> {
|
||||
return ChatModelManager.chatOpenAI.getNumTokens(inputStr);
|
||||
}
|
||||
}
|
||||
81
src/LLMProviders/embeddingManager.ts
Normal file
81
src/LLMProviders/embeddingManager.ts
Normal file
|
|
@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
41
src/LLMProviders/memoryManager.ts
Normal file
41
src/LLMProviders/memoryManager.ts
Normal file
|
|
@ -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();
|
||||
}
|
||||
|
||||
}
|
||||
41
src/LLMProviders/promptManager.ts
Normal file
41
src/LLMProviders/promptManager.ts
Normal file
|
|
@ -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;
|
||||
}
|
||||
}
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
813
src/aiState.ts
813
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<void> {
|
||||
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<void> {
|
||||
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<number> {
|
||||
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<string>(langChainParams.modelDisplayName);
|
||||
const [currentChain, setCurrentChain] = useState<ChainType>(langChainParams.chainType);
|
||||
const [, setChatMemory] = useState<BufferWindowMemory | null>(aiState.memory);
|
||||
const [, setChatMemory] = useState<BufferWindowMemory | null>(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;
|
||||
|
|
|
|||
|
|
@ -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<boolean>;
|
||||
defaultSaveFolder: string;
|
||||
|
|
@ -57,22 +57,20 @@ interface ChatProps {
|
|||
}
|
||||
|
||||
const Chat: React.FC<ChatProps> = ({
|
||||
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<AbortController | null>(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<ChatProps> = ({
|
|||
|
||||
await getAIResponse(
|
||||
userMessage,
|
||||
chatContext,
|
||||
aiState,
|
||||
chainManager,
|
||||
addMessage,
|
||||
setCurrentAiMessage,
|
||||
setAbortController,
|
||||
|
|
@ -173,8 +170,7 @@ const Chat: React.FC<ChatProps> = ({
|
|||
|
||||
await getAIResponse(
|
||||
promptMessageHidden,
|
||||
chatContext,
|
||||
aiState,
|
||||
chainManager,
|
||||
addMessage,
|
||||
setCurrentAiMessage,
|
||||
setAbortController,
|
||||
|
|
@ -203,7 +199,7 @@ const Chat: React.FC<ChatProps> = ({
|
|||
}
|
||||
|
||||
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<ChatProps> = ({
|
|||
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<ChatProps> = ({
|
|||
}
|
||||
|
||||
// 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,
|
||||
|
|
|
|||
|
|
@ -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 {
|
|||
<React.StrictMode>
|
||||
<Chat
|
||||
sharedState={this.sharedState}
|
||||
aiState={this.aiState}
|
||||
chainManager={this.chainManager}
|
||||
emitter={this.emitter}
|
||||
getChatVisibility={this.getChatVisibility}
|
||||
defaultSaveFolder={this.defaultSaveFolder}
|
||||
|
|
|
|||
|
|
@ -11,10 +11,6 @@ export enum ChatModels {
|
|||
GPT_4 = 'gpt-4',
|
||||
GPT_4_TURBO = 'gpt-4-1106-preview',
|
||||
GPT_4_32K = 'gpt-4-32k',
|
||||
// CLAUDE_1 = 'claude-1',
|
||||
// CLAUDE_1_100K = 'claude-1-100k',
|
||||
// CLAUDE_INSTANT_1 = 'claude-instant-1',
|
||||
// CLAUDE_INSTANT_1_100K = 'claude-instant-1-100k',
|
||||
AZURE_GPT_35_TURBO = 'gpt-35-turbo',
|
||||
AZURE_GPT_35_TURBO_16K = 'gpt-35-turbo-16k',
|
||||
GEMINI_PRO = 'gemini-pro',
|
||||
|
|
@ -27,10 +23,6 @@ export enum ChatModelDisplayNames {
|
|||
GPT_4 = 'GPT-4',
|
||||
GPT_4_TURBO = 'GPT-4 TURBO',
|
||||
GPT_4_32K = 'GPT-4 32K',
|
||||
// CLAUDE_1 = 'CLAUDE-1',
|
||||
// CLAUDE_1_100K = 'CLAUDE-1-100K',
|
||||
// CLAUDE_INSTANT_1 = 'CLAUDE-INSTANT',
|
||||
// CLAUDE_INSTANT_1_100K = 'CLAUDE-INSTANT-100K',
|
||||
AZURE_GPT_35_TURBO = 'AZURE GPT-3.5',
|
||||
AZURE_GPT_35_TURBO_16K = 'AZURE GPT-3.5-16K',
|
||||
AZURE_GPT_4 = 'AZURE GPT-4',
|
||||
|
|
@ -57,13 +49,6 @@ export const AZURE_MODELS = new Set([
|
|||
ChatModelDisplayNames.AZURE_GPT_4_32K,
|
||||
]);
|
||||
|
||||
// export const CLAUDE_MODELS = new Set([
|
||||
// ChatModelDisplayNames.CLAUDE_1,
|
||||
// ChatModelDisplayNames.CLAUDE_1_100K,
|
||||
// ChatModelDisplayNames.CLAUDE_INSTANT_1,
|
||||
// ChatModelDisplayNames.CLAUDE_INSTANT_1_100K,
|
||||
// ]);
|
||||
|
||||
export const GOOGLE_MODELS = new Set([
|
||||
ChatModelDisplayNames.GEMINI_PRO,
|
||||
]);
|
||||
|
|
@ -86,10 +71,6 @@ export const DISPLAY_NAME_TO_MODEL: Record<string, string> = {
|
|||
[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<string, string> = {
|
|||
};
|
||||
|
||||
// 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<string, Set<string>> = {
|
||||
[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';
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
12
src/main.ts
12
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<void> | null = null;
|
||||
chatIsVisible = false;
|
||||
dbPrompts: PouchDB.Database;
|
||||
|
|
@ -41,10 +41,10 @@ export default class CopilotPlugin extends Plugin {
|
|||
async onload(): Promise<void> {
|
||||
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<CustomPrompt>('copilot_custom_prompts');
|
||||
this.dbVectorStores = new PouchDB<VectorStoreDocument>('copilot_vector_stores');
|
||||
|
|
@ -447,7 +447,7 @@ export default class CopilotPlugin extends Plugin {
|
|||
.filter(Boolean) as string[];
|
||||
}
|
||||
|
||||
getAIStateParams(): LangChainParams {
|
||||
getChainManagerParams(): LangChainParams {
|
||||
const {
|
||||
openAIApiKey,
|
||||
huggingfaceApiKey,
|
||||
|
|
|
|||
|
|
@ -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<string> => {
|
||||
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);
|
||||
}
|
||||
}
|
||||
};
|
||||
17
src/utils.ts
17
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<string, string> => {
|
||||
const model_to_vendor: Record<string, string> = {};
|
||||
|
||||
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) => {
|
||||
|
|
|
|||
Loading…
Reference in a new issue