update

app.py

@@ -1,110 +1,67 @@
-import chainlit as cl
-from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain.document_loaders.csv_loader import CSVLoader
-from langchain.embeddings import CacheBackedEmbeddings
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.vectorstores import FAISS
-from langchain.chains import RetrievalQA
-from langchain.chat_models import ChatOpenAI
-from langchain.storage import LocalFileStore
-from langchain.prompts.chat import (
-    ChatPromptTemplate,
-    SystemMessagePromptTemplate,
-    HumanMessagePromptTemplate,
+import os
+import openai
+
+from llama_index.query_engine.retriever_query_engine import RetrieverQueryEngine
+from llama_index.callbacks.base import CallbackManager
+from llama_index import (
+    LLMPredictor,
+    ServiceContext,
+    StorageContext,
+    load_index_from_storage,
 )
+from langchain.chat_models import ChatOpenAI
 import chainlit as cl
 
-text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
-
-system_template = """
-Use the following pieces of context to answer the user's question.
-Please respond as if you were Ken from the movie Barbie. Ken is a well-meaning but naive character who loves to Beach. He talks like a typical Californian Beach Bro, but he doesn't use the word "Dude" so much.
-If you don't know the answer, just say that you don't know, don't try to make up an answer.
-You can make inferences based on the context as long as it still faithfully represents the feedback.
-
-Example of your response should be:
 
-```
-The answer is foo
-```
+openai.api_key = os.environ.get("OPENAI_API_KEY")
 
-Begin!
-----------------
-{context}"""
+try:
+    # rebuild storage context
+    storage_context = StorageContext.from_defaults(persist_dir="./storage")
+    # load index
+    index = load_index_from_storage(storage_context)
+except:
+    from llama_index import GPTVectorStoreIndex, SimpleDirectoryReader
 
-messages = [
-    SystemMessagePromptTemplate.from_template(system_template),
-    HumanMessagePromptTemplate.from_template("{question}"),
-]
-prompt = ChatPromptTemplate(messages=messages)
-chain_type_kwargs = {"prompt": prompt}
+    documents = SimpleDirectoryReader("./data").load_data()
+    index = GPTVectorStoreIndex.from_documents(documents)
+    index.storage_context.persist()
 
-@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/")
-    core_embeddings_model = OpenAIEmbeddings()
-    embedder = CacheBackedEmbeddings.from_bytes_store(
-        core_embeddings_model, store, namespace=core_embeddings_model.model
+async def factory():
+    llm_predictor = LLMPredictor(
+        llm=ChatOpenAI(
+            temperature=0,
+            model_name="gpt-3.5-turbo",
+            streaming=True,
+        ),
     )
-    # 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}
+    service_context = ServiceContext.from_defaults(
+        llm_predictor=llm_predictor,
+        chunk_size=512,
+        callback_manager=CallbackManager([cl.LlamaIndexCallbackHandler()]),
     )
 
-    msg.content = f"Index built!"
-    await msg.send()
+    query_engine = index.as_query_engine(
+        service_context=service_context,
+        streaming=True,
+    )
 
-    cl.user_session.set("chain", chain)
+    cl.user_session.set("query_engine", query_engine)
 
 
 @cl.on_message
 async def main(message):
-    chain = cl.user_session.get("chain")
-    cb = cl.AsyncLangchainCallbackHandler(
-        stream_final_answer=False, answer_prefix_tokens=["FINAL", "ANSWER"]
-    )
-    cb.answer_reached = True
-    res = await chain.acall(message, callbacks=[cb], )
-
-    answer = res["result"]
-    source_elements = []
-    visited_sources = set()
+    query_engine = cl.user_session.get("query_engine")  # type: RetrieverQueryEngine
+    response = await cl.make_async(query_engine.query)(message)
 
-    # Get the documents from the user session
-    docs = res["source_documents"]
-    metadatas = [doc.metadata for doc in docs]
-    all_sources = [m["source"] for m in metadatas]
+    response_message = cl.Message(content="")
 
-    for source in all_sources:
-        if source in visited_sources:
-            continue
-        visited_sources.add(source)
-        # Create the text element referenced in the message
-        source_elements.append(
-            cl.Text(content="https://www.imdb.com" + source, name="Review URL")
-        )
+    for token in response.response_gen:
+        await response_message.stream_token(token=token)
 
-    if source_elements:
-        answer += f"\nSources: {', '.join([e.content.decode('utf-8') for e in source_elements])}"
-    else:
-        answer += "\nNo sources found"
+    if response.response_txt:
+        response_message.content = response.response_txt
 
-    await cl.Message(content=answer, elements=source_elements).send()
+    await response_message.send()