fixed generation time bug, for chroma db as well (prev was just model)

documentation_bot/documentation_query.py

@@ -15,10 +15,12 @@
 rag_model_name = "BAAI/bge-small-en"
 generation_model_name = "llama3"  # ollama
 
+
 generation_llm = ChatOllama(
     model=generation_model_name, temperature=0.0
 )
 # Send test message to the generation model
+print("Testing the generation model")
 generation_llm.invoke("test generation")
 
 # Crawl the websites and save the data

documentation_bot/documentation_query_utils.py

@@ -304,6 +304,11 @@ def setup_inference(self, session_id: str) -> None:
         """
         self.store = {}
         self.session_id = session_id
+        self.vectorstore = Chroma(
+            persist_directory=self.chroma_file_path,
+            embedding_function=self.hf_embedding_function,
+        )
+        self.retriever = self.vectorstore.as_retriever(search_kwargs={"k": 1})
 
     def openml_page_search(self, input: str):
 
@@ -312,12 +317,6 @@ def openml_page_search(self, input: str):
         
         """
 
-        vectorstore = Chroma(
-            persist_directory=self.chroma_file_path,
-            embedding_function=self.hf_embedding_function,
-        )
-        retriever = vectorstore.as_retriever(search_kwargs={"k": 1})
-
         ### Contextualize question ###
         contextualize_q_prompt = ChatPromptTemplate.from_messages(
             [
@@ -328,7 +327,7 @@ def openml_page_search(self, input: str):
         )
 
         history_aware_retriever = create_history_aware_retriever(
-            self.generation_llm, retriever, contextualize_q_prompt
+            self.generation_llm, self.retriever, contextualize_q_prompt
         )
 
         ### Answer question ###