Updated code

utils.py

@@ -30,6 +30,9 @@ from langchain.agents.agent_toolkits import create_conversational_retrieval_agen
 from langchain.utilities import SerpAPIWrapper
 from langchain.agents import Tool
 from langchain.agents import load_tools
+from langchain.chat_models import ChatOpenAI
+from langchain.retrievers.multi_query import MultiQueryRetriever
+from langchain.chains import RetrievalQA
 
 load_dotenv()
 
@@ -251,30 +254,70 @@ def load_text_chunks(text_chunks_pkl_dir):
 def load_ensemble_retriver(text_chunks, embeddings, chroma_vectorstore):
     """Load ensemble retiriever with BM25 and Chroma as individual retrievers"""
     bm25_retriever = BM25Retriever.from_documents(text_chunks)
-    bm25_retriever.k = 2
-    chroma_retriever = chroma_vectorstore.as_retriever(search_kwargs={"k": 2})  
+    bm25_retriever.k = 1
+    chroma_retriever = chroma_vectorstore.as_retriever(search_kwargs={"k": 1})  
     ensemble_retriever = EnsembleRetriever(retrievers=[bm25_retriever, chroma_retriever], weights=[0.3, 0.7])
-    return ensemble_retriever
+    retriever_from_llm = MultiQueryRetriever.from_llm(retriever=ensemble_retriever, llm=ChatOpenAI()
+)
+    return retriever_from_llm
 
 
 def load_conversational_retrievel_chain(retriever, llm):
-    '''Load Conversational Retrievel chain,'''
-    _template= """
-    You are a helpful assistant. You do not respond as 'User' or pretend to be 'User'. You only respond once as 'Assistant'.
-    Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question, in its original language.
-    Chat History:
-    {chat_history}
-    Follow Up Input: {question}
-    Standalone question:"""
-
-    CONDENSE_QUESTION_PROMPT = PromptTemplate.from_template(_template)
-    memory = ConversationBufferWindowMemory(return_messages=True,memory_key="chat_history", k=2)
-    conversation_chain = ConversationalRetrievalChain.from_llm(
-        llm=st.session_state["llm"],
-        retriever=st.session_state["ensemble_retriver"],
-        condense_question_prompt=CONDENSE_QUESTION_PROMPT,
-        memory=memory,
-        verbose=True,
+    '''Load Conversational Retrievel agent with following tasks as tools,
+    1) OPM Knowledge base query
+    2) INternet search with SerpAPI
+    This agent combines RAG, chat interfaces, agents.
+    '''
+    # retriever_tool = create_retriever_tool(
+    # retriever, 
+    # "Search_US_Office_of_Personnel_Management_Document",
+    # "Searches and returns documents regarding the U.S. Office of Personnel Management (OPM).")
+    # search_api = SerpAPIWrapper()
+    # search_api_tool = Tool(
+    #         name = "Current_Search",
+    #         func=search_api.run,
+    #         description="useful for when you need to answer questions about current events or the current state of the world"
+    #     )
+    # tools = [retriever_tool]
+    # agent_executor = create_conversational_retrieval_agent(llm, tools, verbose=True, max_token_limit=512)
+    # return agent_executor
+  # string_dialogue = "You are a helpful assistant. You do not respond as 'User' or pretend to be 'User'. You only respond once as 'Assistant'."
+    # _template= """
+    # You are a helpful assistant. You do not respond as 'User' or pretend to be 'User'. You only respond once as 'Assistant'.
+    # Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question, in its original language.
+    # Your answer should in English language only.
+    # Chat History:
+    # {chat_history}
+    # Follow Up Input: {question}
+    # Standalone question:"""
+
+    # CONDENSE_QUESTION_PROMPT = PromptTemplate.from_template(_template)
+    # memory = ConversationBufferMemory(return_messages=True,memory_key="chat_history")
+    # conversation_chain = ConversationalRetrievalChain.from_llm(
+    #     llm=st.session_state["llm"],
+    #     retriever=st.session_state["ensemble_retriver"],
+    #     condense_question_prompt=CONDENSE_QUESTION_PROMPT,
+    #     memory=memory,
+    #     verbose=True,
+    # )
+    template = """You are a helpful assistant. You do not respond as 'User' or pretend to be 'User'. You only respond once as 'Assistant'.
+    Use the following pieces of context to answer the question at the end. If you don't know the answer,\
+    just say that you don't know, don't try to make up an answer.
+
+    {context}
+
+    {history}
+    Question: {question}
+    Helpful Answer:"""
+
+    prompt = PromptTemplate(input_variables=["history", "context", "question"], template=template)
+    memory = ConversationBufferWindowMemory(input_key="question", memory_key="history", k=2)
+
+    qa = RetrievalQA.from_chain_type(
+        llm=llm,
+        chain_type="stuff",
+        retriever=retriever,
+        return_source_documents=True,
+        chain_type_kwargs={"prompt": prompt, "memory": memory},
     )
-    return conversation_chain
-
+    return qa