changing things for better retrieval

app.py

@@ -5,23 +5,16 @@ from langchain_community.embeddings import HuggingFaceBgeEmbeddings
 from langchain_community.llms import Together
 from langchain import hub
 from operator import itemgetter
-from langchain.schema.runnable import RunnableParallel
-from langchain.schema import format_document
-from typing import List, Tuple
 from langchain.chains import LLMChain
-from langchain.chains import RetrievalQA
-from langchain.schema.output_parser import StrOutputParser
-from langchain_community.chat_message_histories import StreamlitChatMessageHistory
-from langchain.memory import ConversationBufferMemory
 from langchain.chains import ConversationalRetrievalChain
-from langchain.memory import ConversationSummaryMemory
 from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder, PromptTemplate
-from langchain.schema.runnable import RunnableLambda, RunnablePassthrough
+from langchain.memory import ConversationBufferMemory
+from langchain_community.chat_message_histories import StreamlitChatMessageHistory
 import time
 
 # Load the embedding function
 model_name = "BAAI/bge-base-en"
-encode_kwargs = {'normalize_embeddings': True} # set True to compute cosine similarity
+encode_kwargs = {'normalize_embeddings': True}
 
 embedding_function = HuggingFaceBgeEmbeddings(
     model_name=model_name,
@@ -47,7 +40,7 @@ llmc = Together(
 )
 
 msgs = StreamlitChatMessageHistory(key="langchain_messages")
-memory = ConversationBufferMemory(chat_memory=msgs)
+memory = ConversationBufferMemory(chat_memory=msgs, memory_key="chat_history", return_messages=True)
 
 DEFAULT_DOCUMENT_PROMPT = PromptTemplate.from_template(template="{page_content}")
 
@@ -60,19 +53,15 @@ def _combine_documents(
 chistory = []
 
 def store_chat_history(role: str, content: str):
-    # Append the new message to the chat history
     chistory.append({"role": role, "content": content})
 
-# Define the Streamlit app
 def app():
     with st.sidebar:
         st.title("dochatter")
-        # Create a dropdown selection box
         option = st.selectbox(
             'Which retriever would you like to use?',
             ('General Medicine', 'RespiratoryFishman', 'RespiratoryMurray', 'MedMRCP2', 'OldMedicine')
         )
-        # Depending on the selected option, choose the appropriate retriever
         if option == 'RespiratoryFishman':
             persist_directory = "./respfishmandbcud/"
             vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding_function, collection_name="fishmannotescud")
@@ -94,37 +83,28 @@ def app():
             vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding_function, collection_name="mrcppassmednotes")
             retriever = vectordb.as_retriever(search_kwargs={"k": 5})
 
-    # Session State
     if "messages" not in st.session_state.keys():
         st.session_state.messages = [{"role": "assistant", "content": "How may I help you?"}]
 
-    _template = """Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question which contains the themes of the conversation. Do not write the question. Do not write the answer.
+    condense_template = """Given the following conversation and a follow-up question, rephrase the follow-up question to be a standalone question which contains the themes of the conversation.
     Chat History:
     {chat_history}
-    Follow Up Input: {question}
+    Follow-Up Input: {question}
     Standalone question:"""
-    CONDENSE_QUESTION_PROMPT = PromptTemplate.from_template(_template)
+    CONDENSE_QUESTION_PROMPT = PromptTemplate.from_template(condense_template)
 
-    template = """You are helping a doctor. Answer with what you know from the context provided. Please be as detailed and thorough. Answer the question based on the following context:
+    answer_template = """You are helping a doctor. Answer with what you know from the context provided. Please be as detailed and thorough. Answer the question based on the following context:
     {context}
-    Question: {question}
-    """
-    ANSWER_PROMPT = ChatPromptTemplate.from_template(template)
-
-    _inputs = RunnableParallel(
-        standalone_question=RunnablePassthrough.assign(
-            chat_history=lambda x: chistory
-        )
-        | CONDENSE_QUESTION_PROMPT
-        | llmc
-        | StrOutputParser(),
+    Question: {question}"""
+    ANSWER_PROMPT = ChatPromptTemplate.from_template(answer_template)
+
+    conversational_qa_chain = ConversationalRetrievalChain(
+        retriever=retriever,
+        memory=memory,
+        combine_docs_chain=_combine_documents,
+        condense_question_chain=LLMChain(llm=llmc, prompt=CONDENSE_QUESTION_PROMPT),
+        qa_chain=LLMChain(llm=llm, prompt=ANSWER_PROMPT)
     )
-    _context = {
-        "context": itemgetter("standalone_question") | retriever | _combine_documents,
-        "question": lambda x: x["standalone_question"],
-    }
-    
-    conversational_qa_chain = _inputs | _context | ANSWER_PROMPT | llm
 
     st.header("Ask Away!")
     for message in st.session_state.messages:
@@ -142,7 +122,7 @@ def app():
     if st.session_state.messages[-1]["role"] != "assistant":
         with st.chat_message("assistant"):
             with st.spinner("Thinking..."):
-                for _ in range(3):  # Retry up to 3 times
+                for _ in range(3):
                     try:
                         response = conversational_qa_chain.invoke(
                             {
@@ -156,7 +136,7 @@ def app():
                         break
                     except Exception as e:
                         st.error(f"An error occurred: {e}")
-                        time.sleep(2)  # Wait 2 seconds before retrying
+                        time.sleep(2)
 
 if __name__ == '__main__':
     app()