Update app.py

app.py

@@ -1,6 +1,6 @@
 import chainlit as cl
 from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain.document_loaders.csv_loader import CSVLoader
+from langchain.document_loaders import WikipediaLoader, CSVLoader
 from langchain.embeddings import CacheBackedEmbeddings
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores import FAISS
@@ -12,7 +12,10 @@ from langchain.prompts.chat import (
     SystemMessagePromptTemplate,
     HumanMessagePromptTemplate,
 )
-import chainlit as cl
+from langchain.retrievers import BM25Retriever, EnsembleRetriever
+from langchain.agents import Tool, ZeroShotAgent, AgentExecutor
+from langchain.agents.agent_toolkits import create_retriever_tool, create_conversational_retrieval_agent
+from langchain import LLMChain
 
 text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
 
@@ -39,40 +42,162 @@ messages = [
 prompt = ChatPromptTemplate(messages=messages)
 chain_type_kwargs = {"prompt": prompt}
 
-@cl.author_rename
-def rename(orig_author: str):
-    rename_dict = {"RetrievalQA": "Consulting The Kens"}
-    return rename_dict.get(orig_author, orig_author)
+# @cl.author_rename
+# def rename(orig_author: str):
+#     rename_dict = {"RetrievalQA": "Consulting The Kens"}
+#     return rename_dict.get(orig_author, orig_author)
 
 @cl.on_chat_start
 async def init():
     msg = cl.Message(content=f"Building Index...")
     await msg.send()
 
-    # build FAISS index from csv
-    loader = CSVLoader(file_path="./data/barbie.csv", source_column="Review_Url")
-    data = loader.load()
-    documents = text_splitter.transform_documents(data)
-    store = LocalFileStore("./cache/")
+    ### start building retrievers, stores and agents
+    llm = ChatOpenAI(model="gpt-3.5-turbo", temperature = 0)
+    
+    barbie_wikipedia_docs = WikipediaLoader(query="Barbie (film)", load_max_docs=1, doc_content_chars_max=1_000_000).load()
+    barbie_csv_docs = CSVLoader(file_path="./barbie_data/barbie.csv", source_column="Review_Url").load()
+    oppenheimer_wikipedia_docs = WikipediaLoader(query="Oppenheimer (film)", load_max_docs=1, doc_content_chars_max=1_000_000).load()
+    oppenheimer_csv_docs = CSVLoader(file_path="./oppenheimer_data/oppenheimer.csv", source_column="Review_Url").load()
+    
+    wikipedia_text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size = 500,
+        chunk_overlap = 0,
+        length_function = len,
+        is_separator_regex= False,
+        separators = ["\n==", "\n", " "] # keep headings, then paragraphs, then sentences
+    )
+    
+    csv_text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size = 1000,
+        chunk_overlap = 50,
+        length_function = len,
+        is_separator_regex= False,
+        separators = ["\n", " "] # keep paragraphs, then sentences
+    )
+    
+    chunked_barbie_wikipedia_docs = wikipedia_text_splitter.transform_documents(barbie_wikipedia_docs)
+    chunked_barbie_csv_docs = csv_text_splitter.transform_documents(barbie_csv_docs)
+    chunked_opp_wikipedia_docs = wikipedia_text_splitter.transform_documents(oppenheimer_wikipedia_docs)
+    chunked_opp_csv_docs = csv_text_splitter.transform_documents(oppenheimer_csv_docs)
+
+    # #### Retrieval and Embedding Strategy
+    # set up cached embeddings store
+    store = LocalFileStore("./shared_cache/")
     core_embeddings_model = OpenAIEmbeddings()
-    embedder = CacheBackedEmbeddings.from_bytes_store(
-        core_embeddings_model, store, namespace=core_embeddings_model.model
+    embedder = CacheBackedEmbeddings.from_bytes_store(core_embeddings_model, store, namespace=core_embeddings_model.model)
+    
+    # We'll implement a `FAISS` vectorstore, and create a retriever from it.
+    barbie_csv_faiss_retriever = await cl.make_async(FAISS.from_documents)(chunked_barbie_csv_docs, embedder).as_retriever()
+    opp_csv_faiss_retriever = await cl.make_async(FAISS.from_documents)(chunked_opp_csv_docs, embedder).as_retriever()
+    opp_wikipedia_faiss_store = await cl.make_async(FAISS.from_documents)(chunked_opp_wikipedia_docs, embedder)
+    opp_wikipedia_faiss_retriever = opp_wikipedia_faiss_store.as_retriever(search_kwargs={"k": 1})
+    
+    # set up BM25 retriever
+    barbie_wikipedia_bm25_retriever = BM25Retriever.from_documents(chunked_barbie_wikipedia_docs)
+    barbie_wikipedia_bm25_retriever.k = 1
+    opp_wikipedia_bm25_retriever = BM25Retriever.from_documents(chunked_opp_wikipedia_docs)
+    opp_wikipedia_bm25_retriever.k = 1
+    
+    # set up FAISS vector store
+    barbie_wikipedia_faiss_store = await cl.make_async(FAISS.from_documents)(chunked_barbie_wikipedia_docs, embedder)
+    barbie_wikipedia_faiss_retriever = barbie_wikipedia_faiss_store.as_retriever(search_kwargs={"k": 1})
+    opp_wikipedia_faiss_store = await cl.make_async(FAISS.from_documents)(chunked_opp_wikipedia_docs, embedder)
+    opp_wikipedia_faiss_retriever = opp_wikipedia_faiss_store.as_retriever(search_kwargs={"k": 1})
+    
+    # set up ensemble retriever
+    barbie_ensemble_retriever = EnsembleRetriever(retrievers=[barbie_wikipedia_bm25_retriever, barbie_wikipedia_faiss_retriever],weights=[0.25, 0.75])
+    opp_ensemble_retriever = EnsembleRetriever(retrievers=[opp_wikipedia_bm25_retriever, opp_wikipedia_faiss_retriever],weights=[0.25, 0.75])
+
+    # #### Retrieval Agent
+    barbie_wikipedia_retrieval_tool = create_retriever_tool(
+        barbie_ensemble_retriever, 
+        "Wikipedia",
+        "Searches and returns documents regarding the plot, history, and cast of the Barbie movie"
+    )
+    
+    barbie_csv_retrieval_tool = create_retriever_tool(
+        barbie_csv_faiss_retriever,
+        "PublicReviews",
+        "Searches and returns documents regarding public reviews of the Barbie movie"
     )
-    # make async docsearch
-    docsearch = await cl.make_async(FAISS.from_documents)(documents, embedder)
-
-    chain = RetrievalQA.from_chain_type(
-        ChatOpenAI(model="gpt-4", temperature=0, streaming=True),
-        chain_type="stuff",
-        return_source_documents=True,
-        retriever=docsearch.as_retriever(),
-        chain_type_kwargs = {"prompt": prompt}
+    
+    barbie_retriever_tools = [barbie_wikipedia_retrieval_tool, barbie_csv_retrieval_tool]
+
+    barbie_retriever_agent_executor = create_conversational_retrieval_agent(llm, barbie_retriever_tools, verbose=True)
+    
+    # #### Multi-source chain
+    system_message = """Use the information from the below two sources to answer any questions.
+    
+    Source 1: public user reviews about the Oppenheimer movie
+    <source1>
+    {source1}
+    </source1>
+    
+    Source 2: the wikipedia page for the Oppenheimer movie including the plot summary, cast, and production information
+    <source2>
+    {source2}
+    </source2>
+    """
+    
+    prompt = ChatPromptTemplate.from_messages([("system", system_message), ("human", "{question}")])
+
+    oppenheimer_multisource_chain = {
+        "source1": (lambda x: x["question"]) | opp_ensemble_retriever,
+        "source2": (lambda x: x["question"]) | opp_csv_faiss_retriever,
+        "question": lambda x: x["question"],
+    } | prompt | llm
+
+    # # Agent Creation
+
+    def query_oppenheimer(input):
+        return oppenheimer_multisource_chain.invoke({"question" : input})
+    
+    tools = [
+        Tool(
+            name = "BarbieInfo",
+            func=barbie_retriever_agent_executor.invoke,
+            description="useful for when you need to answer questions about Barbie. Input should be a fully formed question."
+        ),
+        Tool(
+            name = "OppenheimerInfo",
+            func=query_oppenheimer,
+            description="useful for when you need to answer questions about Oppenheimer. Input should be a fully formed question."
+        ),
+    ]
+
+    prefix = """Have a conversation with a human, answering the following questions as best you can. You have access to the following tools:"""
+    suffix = """Begin!"
+    
+    Question: {input}
+    {agent_scratchpad}"""
+    
+    prompt = ZeroShotAgent.create_prompt(
+        tools, 
+        prefix=prefix, 
+        suffix=suffix, 
+        input_variables=["input", "agent_scratchpad"]
     )
+    
+    llm_chain = LLMChain(llm=llm, prompt=prompt)
+
+    barbenheimer_agent = ZeroShotAgent(llm_chain=llm_chain, tools=tools, verbose=True)
+    barbenheimer_agent_chain = AgentExecutor.from_agent_and_tools(agent=barbenheimer_agent, tools=tools, verbose=True)
+    
+########################################################################################################################################
+    
+    # chain = RetrievalQA.from_chain_type(
+    #     ChatOpenAI(model="gpt-4", temperature=0, streaming=True),
+    #     chain_type="stuff",
+    #     return_source_documents=True,
+    #     retriever=docsearch.as_retriever(),
+    #     chain_type_kwargs = {"prompt": prompt}
+    # )
 
     msg.content = f"Index built!"
     await msg.send()
 
-    cl.user_session.set("chain", chain)
+    cl.user_session.set("chain", barbenheimer_agent_chain)
 
 
 @cl.on_message