Update app.py

app.py

@@ -4,27 +4,26 @@ import streamlit as st
 from langchain.embeddings import HuggingFaceInstructEmbeddings, HuggingFaceEmbeddings
 from langchain.vectorstores.faiss import FAISS
 from huggingface_hub import snapshot_download
-from langchain.chat_models import ChatOpenAI
-from langchain.prompts.chat import (
-    ChatPromptTemplate,
-    SystemMessagePromptTemplate,
-    AIMessagePromptTemplate,
-    HumanMessagePromptTemplate,
-)
-from langchain.schema import (
-    AIMessage,
-    HumanMessage,
-    SystemMessage
+
+from langchain.callbacks import StreamlitCallbackHandler
+from langchain.agents import OpenAIFunctionsAgent, AgentExecutor
+from langchain.agents.agent_toolkits import create_retriever_tool
+from langchain.agents.openai_functions_agent.agent_token_buffer_memory import (
+    AgentTokenBufferMemory,
 )
+from langchain.chat_models import ChatOpenAI
+from langchain.schema import SystemMessage, AIMessage, HumanMessage
+from langchain.prompts import MessagesPlaceholder
+from langsmith import Client
 
-from langchain.chains import ConversationalRetrievalChain
-from langchain.chains.llm import LLMChain
-from langchain.callbacks.base import CallbackManager
-from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
-from langchain.chains.conversational_retrieval.prompts import CONDENSE_QUESTION_PROMPT
-from langchain.chains.question_answering import load_qa_chain
+client = Client()
 
-st.set_page_config(page_title="CFA Level 1", page_icon="📖")
+st.set_page_config(
+    page_title="Chat with CFA Level 1",
+    page_icon="📖",
+    layout="wide",
+    initial_sidebar_state="collapsed",
+)
 
 #Load API Key
 api_key = os.environ["OPENAI_API_KEY"]   
@@ -34,9 +33,10 @@ with st.sidebar:
     book = st.radio("Embedding Model: ",
                    ["Sbert"]
                    )
-  
+
+
 #load embedding models
-@st.experimental_singleton(show_spinner=True)
+@st.cache_resource(show_spinner=True)
 def load_embedding_models(model):
 
     if model == 'Sbert':
@@ -53,14 +53,10 @@ def load_embedding_models(model):
 
     return emb
 
-st.title(f"Talk to CFA Level 1 Book")
-st.markdown("#### Have a conversation with the CFA Curriculum by the CFA Institute 🙊")
-
-
 embeddings = load_embedding_models(book)
 
 ##### functionss ####
-@st.experimental_singleton(show_spinner=False)
+@st.cache_resource(show_spinner=False)
 def load_vectorstore(_embeddings):
     # download from hugging face
     cache_dir="cfa_level_1_cache"
@@ -83,122 +79,75 @@ def load_vectorstore(_embeddings):
             print(target_path)
 
     # load faiss
-    docsearch = FAISS.load_local(folder_path=target_path, embeddings=_embeddings)
-
-    return docsearch
-
-
-@st.experimental_memo(show_spinner=False)
-def load_prompt():
-    system_template="""You are an expert in finance, economics, investing, ethics, derivatives and markets. 
-    Use the following pieces of context to answer the users question. If you don't know the answer, 
-    just say that you don't know, don't try to make up an answer.
-    ALWAYS return a "sources" part in your answer.
-    The "sources" part should be a reference to the source of the context from which you got your answer. List all sources used
-
-    You can use other sources to answer the question if and only if the given context does not have sufficient and relevant information to answer
-    the question.
-    
-    The output should be a markdown code snippet formatted in the following schema:
-    ```json
-    {{
-    answer: is foo
-    sources: xyz
-    }}
-    ```
-    Begin!
-    ----------------
-    {context}"""
-    messages = [
-        SystemMessagePromptTemplate.from_template(system_template),
-        HumanMessagePromptTemplate.from_template("{question}")
-    ]
-    prompt = ChatPromptTemplate.from_messages(messages)
-
-    return prompt
-
-    
-@st.experimental_singleton(show_spinner=False)
-def load_chain():
-    '''Load langchain Conversational Retrieval Chain'''
-    
-    vectorstore = load_vectorstore(embeddings)
-    llm = ChatOpenAI(temperature=0, model_name='gpt-4-32k-0613')
-    streaming_llm = ChatOpenAI(model_name='gpt-4-32k-0613',
-                               streaming=True, 
-                               callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]), 
-                               verbose=True, 
-                               temperature=0)
-
-    question_generator = LLMChain(llm=llm, prompt=CONDENSE_QUESTION_PROMPT)
-    doc_chain = load_qa_chain(streaming_llm, chain_type="stuff", prompt=load_prompt())
-
-    qa = ConversationalRetrievalChain(
-        retriever=vectorstore.as_retriever(), 
-        combine_docs_chain=doc_chain, 
-        question_generator=question_generator,
-        return_source_documents=True)
-    
-    return qa
-
-chat_history = []
-
-def get_answer(question):
-    '''Generate an answer from the chain'''
-
-    chain = load_chain()
-    result = chain({"question": question, "chat_history": chat_history})
-
-    answer = result["answer"]
-    
-    # pages
-    unique_sources = set()
-    for item in result['source_documents']:
-        unique_sources.add(item.metadata['source'].split(',')[1])
-
-    unique_pages = ""
-    for item in unique_sources:
-        unique_pages += str(item) + ", "
-
-    # will look like 1, 2, 3,
-    pages = unique_pages[:-2] # removes the last comma and space
-
-    # source text
-    full_source = ""
-    for item in result['source_documents']:
-        full_source += f"- **{item.metadata['source']}**" + "\n" + item.page_content + "\n\n"
-
-    # will look like:
-    # - Page: {number}
-    #  {extracted text from book}
-    extract = full_source
-
-    return answer, pages, extract
-
-
-##### main ####
-user_input = st.text_area("Your question", "What is an MBS and who are the main issuers and investors of the MBS market?", key="input")
-
-col1, col2 = st.columns([10, 1])
-
-# show question
-col1.write(f"**You:** {user_input}")
-
-# ask button to the right of the displayed question
-ask = col2.button("Ask", type="primary")
-
-
-if ask:
-
-    with st.spinner("this can take about a minute for your first question because some models have to be downloaded 🥺👉🏻👈🏻"):
-        try:
-            answer, pages, extract = get_answer(question=user_input)
-        except Exception as e:
-            st.write(f"Error with Download: {e}")
-            st.stop()
-
-    st.write(f"{answer}")
-
-    # sources
-    with st.expander(label = f"From: {pages}", expanded = False):
-        st.markdown(extract)
+    vectorstore = FAISS.load_local(folder_path=target_path, embeddings=_embeddings)
+
+    return vectorstore.as_retriever(search_kwargs={"k": 4})
+
+tool = create_retriever_tool(
+    load_vectorstore(),
+    "search_cfa_docs",
+    "Searches and returns documents regarding the CFA level 1 curriculum. CFA is a rigorous program for investment professionals which covers topics such as ethics, corporate finance, economics, fixed income, equities and derivatives markets. You do not know anything about the CFA program, so if you are ever asked about CFA material or curriculum you should use this tool.",
+)
+tools = [tool]
+llm = ChatOpenAI(temperature=0, streaming=True, model="gpt-4")
+message = SystemMessage(
+    content=(
+        "You are a helpful chatbot who is tasked with answering questions about the CFA level 1 program. "
+        "Unless otherwise explicitly stated, it is probably fair to assume that questions are about the CFA program and materials. "
+        "If there is any ambiguity, you probably assume they are about that."
+    )
+)
+
+prompt = OpenAIFunctionsAgent.create_prompt(
+    system_message=message,
+    extra_prompt_messages=[MessagesPlaceholder(variable_name="history")],
+)
+agent = OpenAIFunctionsAgent(llm=llm, tools=tools, prompt=prompt)
+agent_executor = AgentExecutor(
+    agent=agent,
+    tools=tools,
+    verbose=True,
+    return_intermediate_steps=True,
+)
+memory = AgentTokenBufferMemory(llm=llm)
+starter_message = "Ask me anything about the CFA Level 1 Curriculum!"
+if "messages" not in st.session_state or st.sidebar.button("Clear message history"):
+    st.session_state["messages"] = [AIMessage(content=starter_message)]
+
+
+def send_feedback(run_id, score):
+    client.create_feedback(run_id, "user_score", score=score)
+
+
+for msg in st.session_state.messages:
+    if isinstance(msg, AIMessage):
+        st.chat_message("assistant").write(msg.content)
+    elif isinstance(msg, HumanMessage):
+        st.chat_message("user").write(msg.content)
+    memory.chat_memory.add_message(msg)
+
+
+if prompt := st.chat_input(placeholder=starter_message):
+    st.chat_message("user").write(prompt)
+    with st.chat_message("assistant"):
+        st_callback = StreamlitCallbackHandler(st.container())
+        response = agent_executor(
+            {"input": prompt, "history": st.session_state.messages},
+            callbacks=[st_callback],
+            include_run_info=True,
+        )
+        st.session_state.messages.append(AIMessage(content=response["output"]))
+        st.write(response["output"])
+        memory.save_context({"input": prompt}, response)
+        st.session_state["messages"] = memory.buffer
+        run_id = response["__run"].run_id
+
+        col_blank, col_text, col1, col2 = st.columns([10, 2, 1, 1])
+        with col_text:
+            st.text("Feedback:")
+
+        with col1:
+            st.button("👍", on_click=send_feedback, args=(run_id, 1))
+
+        with col2:
+            st.button("👎", on_click=send_feedback, args=(run_id, 0))