app.py

import streamlit as st
import os
from langchain_community.vectorstores import Chroma
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
import time

# Load the embedding function
model_name = "BAAI/bge-base-en"
encode_kwargs = {'normalize_embeddings': True} # set True to compute cosine similarity

embedding_function = HuggingFaceBgeEmbeddings(
    model_name=model_name,
    encode_kwargs=encode_kwargs
)

# Load the LLM
llm = Together(
    model="mistralai/Mixtral-8x22B-Instruct-v0.1",
    temperature=0.2,
    max_tokens=19096,
    top_k=10,
    together_api_key=os.environ['pilotikval']
)

# Load the summarizeLLM
llmc = Together(
    model="mistralai/Mixtral-8x7B-Instruct-v0.1",
    temperature=0.2,
    max_tokens=1024,
    top_k=1,
    together_api_key=os.environ['pilotikval']
)

msgs = StreamlitChatMessageHistory(key="langchain_messages")
memory = ConversationBufferMemory(chat_memory=msgs)

DEFAULT_DOCUMENT_PROMPT = PromptTemplate.from_template(template="{page_content}")

def _combine_documents(
        docs, document_prompt=DEFAULT_DOCUMENT_PROMPT, document_separator="\n\n"
    ):
    doc_strings = [format_document(doc, document_prompt) for doc in docs]
    return document_separator.join(doc_strings)

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")
            retriever = vectordb.as_retriever(search_kwargs={"k": 5})
        elif option == 'RespiratoryMurray':
            persist_directory = "./respmurray/"
            vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding_function, collection_name="respmurraynotes")
            retriever = vectordb.as_retriever(search_kwargs={"k": 5})
        elif option == 'MedMRCP2':
            persist_directory = "./medmrcp2store/"
            vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding_function, collection_name="medmrcp2notes")
            retriever = vectordb.as_retriever(search_kwargs={"k": 5})
        elif option == 'General Medicine':
            persist_directory = "./oxfordmedbookdir/"
            vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding_function, collection_name="oxfordmed")
            retriever = vectordb.as_retriever(search_kwargs={"k": 7})
        else:
            persist_directory = "./mrcpchromadb/"
            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.
    Chat History:
    {chat_history}
    Follow Up Input: {question}
    Standalone question:"""
    CONDENSE_QUESTION_PROMPT = PromptTemplate.from_template(_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:
    {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(),
    )
    _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:
        with st.chat_message(message["role"]):
            st.write(message["content"])
            store_chat_history(message["role"], message["content"])

    prompts2 = st.chat_input("Say something")

    if prompts2:
        st.session_state.messages.append({"role": "user", "content": prompts2})
        with st.chat_message("user"):
            st.write(prompts2)

    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
                    try:
                        response = conversational_qa_chain.invoke(
                            {
                                "question": prompts2,
                                "chat_history": chistory,
                            }
                        )
                        st.write(response)
                        message = {"role": "assistant", "content": response}
                        st.session_state.messages.append(message)
                        break
                    except Exception as e:
                        st.error(f"An error occurred: {e}")
                        time.sleep(2)  # Wait 2 seconds before retrying

if __name__ == '__main__':
    app()