app.py

import streamlit as st
import os
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceBgeEmbeddings
from langchain.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.memory 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


# 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 ChromaDB vector store
# persist_directory="./mrcpchromadb/"
# vectordb = Chroma(persist_directory=persist_directory, embedding_function=embedding_function,collection_name="mrcppassmednotes")




# Load the LLM
llm = Together(
    model="mistralai/Mixtral-8x7B-Instruct-v0.1",
    temperature=0.2,
    max_tokens=4096,
    top_k=4,
    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?',
            ('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})
            retriever = retriever # replace with your actual retriever
        
        if 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})
            retriever = retriever

        if 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})
            retriever = retriever
        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})
            retriever = retriever # replace with your actual retriever
            retriever = retriever  # replace with your actual retriever

        #template = """You are an AI chatbot having a conversation with a human. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.
        #{context}
        #{history}
        #Human: {human_input}
        #AI: """
        #prompt = PromptTemplate(input_variables=["history", "question"], template=template)
        #template = st.text_area("Template", value=template, height=180)
        #prompt2 = ChatPromptTemplate.from_template(template)

    
    

    # Session State
    # Store LLM generated responses
    if "messages" not in st.session_state.keys():
        st.session_state.messages = [{"role": "assistant", "content": "How may I help you?"}]
    

    

    


    


    ## Retry lets go

    _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!")
    # Display the messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.write(message["content"])
            store_chat_history(message["role"], message["content"])

    # prompt = hub.pull("rlm/rag-prompt")
    
    
    
    
    prompts2 = st.chat_input("Say something")

    # Implement using different book sources, if statements
    

    

    
    

    

    

    
    
    
    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..."):
                response = conversational_qa_chain.invoke(
                {
                    "question": prompts2,
                    "chat_history": chistory,
                }
            )
                st.write(response)
        message = {"role": "assistant", "content": response}
        st.session_state.messages.append(message)



    
        
        
        
        


   

    # Create a button to submit the question
    


    

        
        

        
        
        

        
        

        

    
# Initialize history
history = []

if __name__ == '__main__':
    app()