Updated code

app.py

@@ -58,7 +58,7 @@ if "ensemble_retriver" not in st.session_state:
     st.session_state["ensemble_retriver"] = load_ensemble_retriver(text_chunks=st.session_state["text_chunks"], embeddings=st.session_state["embeddings"], chroma_vectorstore=st.session_state["vector_db"] )
 
 if "conversation_chain" not in st.session_state:
-    st.session_state["conversation_chain"] = load_conversational_retrievel_chain(retriever=st.session_state["ensemble_retriver"], llm=st.session_state["llm"])
+    st.session_state["conversation_chain"] = load_conversational_retrievel_chain(retriever=st.session_state["ensemble_retriver"], llm=st.session_state["llm"])    
 
 
 
@@ -193,8 +193,9 @@ if st.session_state["vector_db"] and st.session_state["llm"]:
 
 
     def generate_llm_response(conversation_chain, prompt_input):
-        output= conversation_chain({'question': prompt_input})
-        return output['answer']
+        # output= conversation_chain({'question': prompt_input})
+        res = conversation_chain(prompt_input)
+        return res['result']
 
 
     # User-provided prompt
@@ -214,9 +215,13 @@ if st.session_state["vector_db"] and st.session_state["llm"]:
                 for item in response:
                     full_response += item
                 placeholder.markdown(full_response)
+                # The following logic will work in the way given below.
+                # -- Check if intermediary steps are present in the output of the given prompt.
+                # -- If not, we can conclude that, agent has used internet search as tool.
+                # -- Check if intermediary steps are present in the output of the prompt.
+                # -- If intermediary steps are present, it means agent has used exising custom knowledge base for iformation retrival and therefore we need to give souce docs as output along with LLM's reponse.
                 if response:
                     st.text("-------------------------------------")
-                    #Getting source docs
                     docs= st.session_state["ensemble_retriver"].get_relevant_documents(prompt)   
                     source_doc_list= []  
                     for doc in docs:
@@ -233,6 +238,9 @@ if st.session_state["vector_db"] and st.session_state["llm"]:
                             st.write("---")  # Add a separator between entries
                     message = {"role": "assistant", "content": full_response, "Source":merged_source_doc}
                     st.session_state.messages.append(message)
-
+                # else:
+                    # with st.expander("source"):
+                    #     message = {"role": "assistant", "content": full_response, "Source":""}
+                    #     st.session_state.messages.append(message)
         end = timeit.default_timer()
         print(f"Time to retrieve response: {end - start}")

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 = 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,
+    '''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 = ConversationBufferMemory(input_key="question", memory_key="history")
+
+    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