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

@@ -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<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...');
+    this.vectorStore = await MemoryVectorStore.fromDocuments(
+      docs, embeddingsAPI,
+    );
+    ChainFactory.setVectorStore(this.vectorStore, docHash);
+    console.log('Vector store created successfully.');
+  }
+
   async countTokens(inputStr: string): Promise<number> {
     return AIState.chatOpenAI.getNumTokens(inputStr);
   }

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<string, BaseChain> = new Map();
+  public static vectorStoreMap: Map<string, VectorStore> = 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

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<ChatProps> = ({
     }
     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 = () => {

src/components/ChatComponents/ChatIcons.tsx

@@ -78,7 +78,7 @@ const ChatIcons: React.FC<ChatIconsProps> = ({
 
       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<ChatIconsProps> = ({
       </button>
       <button className='chat-icon-button' onClick={onUseActiveNoteAsContext}>
         <UseActiveNoteAsContextIcon className='icon-scaler' />
-        <span className="tooltip-text">QA: Active Short Note</span>
+        <span className="tooltip-text">Use Active Note as Context</span>
       </button>
       <div className="chat-icon-selection-tooltip">
         <div className="select-wrapper">
@@ -130,7 +130,7 @@ const ChatIcons: React.FC<ChatIconsProps> = ({
             onChange={handleChainChange}
           >
             <option value='llm_chain'>Conversation</option>
-            <option value='retrieval_qa'>QA: Long Note</option>
+            <option value='retrieval_qa'>QA: Active Note</option>
           </select>
           <span className="tooltip-text">Chain Selection<br/>(clears history!)</span>
         </div>