Update app.py

app.py

@@ -1,80 +1,105 @@
-# You can find this code for Chainlit python streaming here (https://docs.chainlit.io/concepts/streaming/python)
-
-# OpenAI Chat completion
 import os
-from openai import AsyncOpenAI  # importing openai for API usage
-import chainlit as cl  # importing chainlit for our app
-from chainlit.prompt import Prompt, PromptMessage  # importing prompt tools
-from chainlit.playground.providers import ChatOpenAI  # importing ChatOpenAI tools
-from dotenv import load_dotenv
+from typing import List
 
-api_key = os.getenv("OPENAI_API_KEY")
+from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.vectorstores import Chroma
+from langchain.chains import (
+    ConversationalRetrievalChain,
+)
+from langchain.chat_models import ChatOpenAI
 
-# ChatOpenAI Templates
-system_template = """You are a helpful assistant who always speaks in a pleasant tone!
-"""
+from langchain.docstore.document import Document
+from langchain.memory import ChatMessageHistory, ConversationBufferMemory
 
-user_template = """{input}
-Think through your response step by step.
-"""
+import chainlit as cl
 
+os.environ["OPENAI_API_KEY"] = "OPENAI_API_KEY"
 
-@cl.on_chat_start  # marks a function that will be executed at the start of a user session
-async def start_chat():
-    settings = {
-        "model": "gpt-3.5-turbo",
-        "temperature": 0,
-        "max_tokens": 500,
-        "top_p": 1,
-        "frequency_penalty": 0,
-        "presence_penalty": 0,
-    }
+text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
 
-    cl.user_session.set("settings", settings)
 
+@cl.on_chat_start
+async def on_chat_start():
+    files = None
 
-@cl.on_message  # marks a function that should be run each time the chatbot receives a message from a user
-async def main(message: cl.Message):
-    settings = cl.user_session.get("settings")
-
-    client = AsyncOpenAI()
-
-    print(message.content)
-
-    prompt = Prompt(
-        provider=ChatOpenAI.id,
-        messages=[
-            PromptMessage(
-                role="system",
-                template=system_template,
-                formatted=system_template,
-            ),
-            PromptMessage(
-                role="user",
-                template=user_template,
-                formatted=user_template.format(input=message.content),
-            ),
-        ],
-        inputs={"input": message.content},
-        settings=settings,
+    # Wait for the user to upload a file
+    while files == None:
+        files = await cl.AskFileMessage(
+            content="Please upload a text file to begin!",
+            accept=["text/plain"],
+            max_size_mb=20,
+            timeout=180,
+        ).send()
+
+    file = files[0]
+
+    msg = cl.Message(content=f"Processing `{file.name}`...")
+    await msg.send()
+
+    with open(file.path, "r", encoding="utf-8") as f:
+        text = f.read()
+
+    # Split the text into chunks
+    texts = text_splitter.split_text(text)
+
+    # Create a metadata for each chunk
+    metadatas = [{"source": f"{i}-pl"} for i in range(len(texts))]
+
+    # Create a Chroma vector store
+    embeddings = OpenAIEmbeddings()
+    docsearch = await cl.make_async(Chroma.from_texts)(
+        texts, embeddings, metadatas=metadatas
+    )
+
+    message_history = ChatMessageHistory()
+
+    memory = ConversationBufferMemory(
+        memory_key="chat_history",
+        output_key="answer",
+        chat_memory=message_history,
+        return_messages=True,
     )
 
-    print([m.to_openai() for m in prompt.messages])
+    # Create a chain that uses the Chroma vector store
+    chain = ConversationalRetrievalChain.from_llm(
+        ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0, streaming=True),
+        chain_type="stuff",
+        retriever=docsearch.as_retriever(),
+        memory=memory,
+        return_source_documents=True,
+    )
 
-    msg = cl.Message(content="")
+    # Let the user know that the system is ready
+    msg.content = f"Processing `{file.name}` done. You can now ask questions!"
+    await msg.update()
 
-    # Call OpenAI
-    async for stream_resp in await client.chat.completions.create(
-        messages=[m.to_openai() for m in prompt.messages], stream=True, **settings
-    ):
-        token = stream_resp.choices[0].delta.content
-        if not token:
-            token = ""
-        await msg.stream_token(token)
+    cl.user_session.set("chain", chain)
 
-    # Update the prompt object with the completion
-    prompt.completion = msg.content
-    msg.prompt = prompt
 
-    # Send and close the message stream
-    await msg.send()
+@cl.on_message
+async def main(message: cl.Message):
+    chain = cl.user_session.get("chain")  # type: ConversationalRetrievalChain
+    cb = cl.AsyncLangchainCallbackHandler()
+
+    res = await chain.acall(message.content, callbacks=[cb])
+    answer = res["answer"]
+    source_documents = res["source_documents"]  # type: List[Document]
+
+    text_elements = []  # type: List[cl.Text]
+
+    if source_documents:
+        for source_idx, source_doc in enumerate(source_documents):
+            source_name = f"source_{source_idx}"
+            # Create the text element referenced in the message
+            text_elements.append(
+                cl.Text(content=source_doc.page_content, name=source_name, display="side")
+            )
+        source_names = [text_el.name for text_el in text_elements]
+
+        if source_names:
+            answer += f"\nSources: {', '.join(source_names)}"
+        else:
+            answer += "\nNo sources found"
+
+    await cl.Message(content=answer, elements=text_elements).send()