From b4e25ac7de7adea9f70500375cd2abdd2b5fcbb5 Mon Sep 17 00:00:00 2001 From: Logan Yang Date: Sun, 28 May 2023 21:05:02 -0700 Subject: [PATCH] Transition to unlimited context QA (#40) * Transition to unlimited context qa * Fix issue where switching active note does not switch the context for QA --- src/aiState.ts | 58 ++++++++++++--------- src/chainFactory.ts | 41 +++------------ src/components/Chat.tsx | 47 +++++++++-------- src/components/ChatComponents/ChatIcons.tsx | 6 +-- 4 files changed, 68 insertions(+), 84 deletions(-) diff --git a/src/aiState.ts b/src/aiState.ts index 2881f32c..a28f8c37 100644 --- a/src/aiState.ts +++ b/src/aiState.ts @@ -17,6 +17,7 @@ import { RetrievalQAChain, } from "langchain/chains"; import { ChatOpenAI } from 'langchain/chat_models/openai'; +import { VectorStore } from 'langchain/dist/vectorstores/base'; import { HuggingFaceInferenceEmbeddings } from "langchain/embeddings/hf"; import { OpenAIEmbeddings } from "langchain/embeddings/openai"; import { BufferWindowMemory } from "langchain/memory"; @@ -57,6 +58,7 @@ class AIState { memory: BufferWindowMemory; langChainParams: LangChainParams; chatPrompt: ChatPromptTemplate; + vectorStore: VectorStore constructor(langChainParams: LangChainParams) { this.langChainParams = langChainParams; @@ -155,40 +157,29 @@ class AIState { } const docHash = ChainFactory.getDocumentHash(options.noteContent); - if (ChainFactory.instances.has(docHash)) { - AIState.retrievalChain = ChainFactory.getRetrievalChain(docHash); - this.langChainParams.chainType = RETRIEVAL_QA_CHAIN; - console.log('Set chain:', RETRIEVAL_QA_CHAIN); - return; + const vectorStore = ChainFactory.vectorStoreMap.get(docHash); + if (vectorStore) { + AIState.retrievalChain = RetrievalQAChain.fromLLM( + AIState.chatOpenAI, + vectorStore.asRetriever(), + ); + console.log('Existing vector store for document hash: ', docHash); } else { - // Should not create new a new chain and index if there's already one - const textSplitter = new RecursiveCharacterTextSplitter({ chunkSize: 1000 }); - - const docs = await textSplitter.createDocuments([options.noteContent]); - const embeddingsAPI = this.getEmbeddingsAPI(); - try { - // Note: HF can give 503 errors frequently (it's free) - console.log('Creating vector store...'); - const vectorStore = await MemoryVectorStore.fromDocuments( - docs, embeddingsAPI, + await this.buildIndex(options.noteContent, docHash); + AIState.retrievalChain = RetrievalQAChain.fromLLM( + AIState.chatOpenAI, + this.vectorStore.asRetriever(), ); - console.log('Vector store created successfully.'); - /* Create or retrieve the chain */ - AIState.retrievalChain = ChainFactory.createRetrievalChain( - { - llm: AIState.chatOpenAI, - retriever: vectorStore.asRetriever(), - }, - docHash, - ); - this.langChainParams.chainType = RETRIEVAL_QA_CHAIN; - console.log('Set chain:', RETRIEVAL_QA_CHAIN); + console.log('New retrieval qa chain created for document hash: ', docHash); } catch (error) { new Notice('Failed to create vector store, please try again:', error); console.log('Failed to create vector store, please try again.:', error); } } + + this.langChainParams.chainType = RETRIEVAL_QA_CHAIN; + console.log('Set chain:', RETRIEVAL_QA_CHAIN); break; } default: @@ -197,6 +188,21 @@ class AIState { } } + 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...'); + this.vectorStore = await MemoryVectorStore.fromDocuments( + docs, embeddingsAPI, + ); + ChainFactory.setVectorStore(this.vectorStore, docHash); + console.log('Vector store created successfully.'); + } + async countTokens(inputStr: string): Promise { return AIState.chatOpenAI.getNumTokens(inputStr); } diff --git a/src/chainFactory.ts b/src/chainFactory.ts index a480426f..dccc66cd 100644 --- a/src/chainFactory.ts +++ b/src/chainFactory.ts @@ -4,9 +4,9 @@ import { BaseChain, ConversationChain, ConversationalRetrievalQAChain, - LLMChainInput, - RetrievalQAChain, + LLMChainInput } from "langchain/chains"; +import { VectorStore } from 'langchain/dist/vectorstores/base'; import { BaseRetriever } from "langchain/schema"; @@ -44,6 +44,7 @@ export const SUPPORTED_CHAIN_TYPES = new Set([ class ChainFactory { public static instances: Map = new Map(); + public static vectorStoreMap: Map = new Map(); public static getLLMChain(args: LLMChainInput): BaseChain { let instance = ChainFactory.instances.get(LLM_CHAIN); @@ -59,41 +60,15 @@ class ChainFactory { return MD5(sourceDocument).toString(); } - /** - * Get the retrieval chain for a given source document. If the document exists - * and is not changed, it retrieves its previous chain instance. Otherwise, - * it creates a new chain and put it into the chain factory map where the key - * is the full document's MD5 hash. - */ - public static getRetrievalChain( - inputDocHash: string, - ): RetrievalQAChain { - let instance; - if (ChainFactory.instances.has(inputDocHash)) { - // Use the existing chain when the note has been indexed - instance = ChainFactory.instances.get(inputDocHash) as RetrievalQAChain; - console.log('Retrieval qa chain retrieved for document hash: ', inputDocHash); - } - - return instance as RetrievalQAChain; + public static setVectorStore(vectorStore: VectorStore, docHash: string): void { + ChainFactory.vectorStoreMap.set(docHash, vectorStore); } - public static createRetrievalChain( - args: RetrievalChainParams, - inputDocHash: string, - ): RetrievalQAChain { - // Create a new retrieval chain when the note hasn't been indexed - const argsRetrieval = args as RetrievalChainParams; - const instance = RetrievalQAChain.fromLLM( - argsRetrieval.llm, argsRetrieval.retriever, argsRetrieval.options - ); - ChainFactory.instances.set(inputDocHash, instance); - console.log('New retrieval qa chain created for document hash: ', inputDocHash); - - return instance as RetrievalQAChain; + public static getVectorStore(docHash: string): VectorStore | undefined { + return ChainFactory.vectorStoreMap.get(docHash); } - public static getConversationalRetrievalChain( + public static createConversationalRetrievalChain( args: ConversationalRetrievalChainParams ): ConversationalRetrievalQAChain { // Create a new retrieval chain every time, not singleton diff --git a/src/components/Chat.tsx b/src/components/Chat.tsx index 076ceb95..62960368 100644 --- a/src/components/Chat.tsx +++ b/src/components/Chat.tsx @@ -1,4 +1,5 @@ import AIState, { useAIState } from '@/aiState'; +import ChainFactory from '@/chainFactory'; import ChatIcons from '@/components/ChatComponents/ChatIcons'; import ChatInput from '@/components/ChatComponents/ChatInput'; import ChatMessages from '@/components/ChatComponents/ChatMessages'; @@ -27,8 +28,7 @@ import { rewriteTweetThreadSelectionPrompt, simplifyPrompt, summarizePrompt, - tocPrompt, - useNoteAsContextPrompt, + tocPrompt } from '@/utils'; import { EventEmitter } from 'events'; import { Notice, TFile } from 'obsidian'; @@ -130,28 +130,31 @@ const Chat: React.FC = ({ } const noteContent = await getFileContent(file); const noteName = getFileName(file); + if (!noteContent) { + new Notice('No note content found.'); + console.error('No note content found.'); + return; + } - // Set the context based on the noteContent - const prompt = useNoteAsContextPrompt(noteName, noteContent); + let activeNoteOnMessage: ChatMessage; + const docHash = ChainFactory.getDocumentHash(noteContent); + const vectorStore = ChainFactory.vectorStoreMap.get(docHash); + if (vectorStore) { + activeNoteOnMessage = { + sender: AI_SENDER, + message: `I have Read [[${noteName}]].\n\n Please switch to "QA: Active Note" to ask questions about it.`, + isVisible: true, + }; + } else { + await aiState.buildIndex(noteContent, docHash); + activeNoteOnMessage = { + sender: AI_SENDER, + message: `Reading [[${noteName}]]...\n\n Please switch to "QA: Active Note" to ask questions about it.`, + isVisible: true, + }; + } - // Send the prompt as a user message - const promptMessage: ChatMessage = { - sender: USER_SENDER, - message: prompt, - isVisible: false, - }; - addMessage(promptMessage); - - // Hide the prompt from the user - await getAIResponse( - promptMessage, - chatContext, - aiState, - addMessage, - setCurrentAiMessage, - setAbortController, - debug, - ); + addMessage(activeNoteOnMessage); }; const clearCurrentAiMessage = () => { diff --git a/src/components/ChatComponents/ChatIcons.tsx b/src/components/ChatComponents/ChatIcons.tsx index fb4d2508..af1d31a3 100644 --- a/src/components/ChatComponents/ChatIcons.tsx +++ b/src/components/ChatComponents/ChatIcons.tsx @@ -78,7 +78,7 @@ const ChatIcons: React.FC = ({ const activeNoteOnMessage: ChatMessage = { sender: AI_SENDER, - message: `OK I've read this long note [[${noteName}]]. Feel free to ask me questions about it.`, + message: `OK Feel free to ask me questions about [[${noteName}]].`, isVisible: true, }; addMessage(activeNoteOnMessage); @@ -119,7 +119,7 @@ const ChatIcons: React.FC = ({
@@ -130,7 +130,7 @@ const ChatIcons: React.FC = ({ onChange={handleChainChange} > - + Chain Selection
(clears history!)