Create app.py

app.py

@@ -0,0 +1,139 @@
+import os
+from typing import List
+import PyPDF2
+from io import BytesIO
+
+from langchain_openai import OpenAIEmbeddings
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import Chroma
+from langchain.chains import ConversationalRetrievalChain
+from langchain_openai import ChatOpenAI
+
+from langchain.docstore.document import Document
+from langchain.memory import ChatMessageHistory, ConversationBufferMemory
+
+import chainlit as cl
+
+open_aikey = os.getenv('openaikey')
+
+OPENAI_API_KEY = open_aikey
+
+text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
+
+template_questions = [
+    "1. What is this document about?",
+    "2. What company is this document about?",
+]
+
+@cl.on_chat_start
+async def on_chat_start():
+    files = None
+
+    # Wait for the user to upload a file
+    while files is None:
+        files = await cl.AskFileMessage(
+            content="Please upload a text or PDF file to begin!",
+            accept=["text/plain", "application/pdf"],
+            max_size_mb=20,
+            timeout=180,
+        ).send()
+
+    file = files[0]
+
+    msg = cl.Message(content=f"Processing `{file.name}`...")
+    await msg.send()
+
+    # Check file type and read accordingly
+    if file.type == "text/plain":
+        with open(file.path, "r", encoding="utf-8") as f:
+            text = f.read()
+    elif file.type == "application/pdf":
+        # Read PDF from a temporary path given by file.path
+        with open(file.path, "rb") as f:
+            pdf_reader = PyPDF2.PdfReader(f)
+            text = ""
+            for page in pdf_reader.pages:
+                text += page.extract_text() or ""
+
+    # Process the text normally from here
+    texts = text_splitter.split_text(text)
+    metadatas = [{"source": f"{i}-pl"} for i in range(len(texts))]
+    embeddings = OpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)
+    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,
+    )
+
+    chain = ConversationalRetrievalChain.from_llm(
+        ChatOpenAI(model_name="gpt-4o-mini", temperature=0, streaming=True, openai_api_key=OPENAI_API_KEY),
+        chain_type="stuff",
+        retriever=docsearch.as_retriever(),
+        memory=memory,
+        return_source_documents=True,
+    )
+
+    cl.user_session.set("chain", chain)
+
+    # Answer the template questions
+    answers = []
+    for question in template_questions:
+        res = await chain.acall(question)
+        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)
+                )
+            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"
+
+        answers.append(answer)
+
+    for i, question in enumerate(template_questions):
+        await cl.Message(content=f"**{question}**\n{answers[i]}").send()
+
+    msg.content = f"Processing `{file.name}` done. You can now ask more questions!"
+    await msg.update()
+
+@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)
+            )
+        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()