Spaces:
Sleeping
Sleeping
File size: 4,724 Bytes
5a9839d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 |
import os
from chainlit.types import AskFileResponse
from utilities_2.openai_utils.prompts import (
UserRolePrompt,
SystemRolePrompt,
AssistantRolePrompt,
)
from utilities_2.openai_utils.embedding import EmbeddingModel
from utilities_2.vectordatabase import VectorDatabase
from utilities_2.openai_utils.chatmodel import ChatOpenAI
import chainlit as cl
from utilities.text_utils import FileLoader
from utilities.pipeline import RetrievalAugmentedQAPipeline
# from utilities.vector_database import QdrantDatabase
def process_file(file, use_rct):
fileLoader = FileLoader()
return fileLoader.load_file(file, use_rct)
system_template = """\
Use the following context to answer a users question.
If you cannot find the answer in the context, say you don't know the answer.
The context contains the text from a document. Refer to it as the document not the context.
"""
system_role_prompt = SystemRolePrompt(system_template)
user_prompt_template = """\
Context:
{context}
Question:
{question}
"""
user_role_prompt = UserRolePrompt(user_prompt_template)
@cl.on_chat_start
async def on_chat_start():
# get user inputs
res = await cl.AskActionMessage(
content="Do you want to use Qdrant?",
actions=[
cl.Action(name="yes", value="yes", label="✅ Yes"),
cl.Action(name="no", value="no", label="❌ No"),
],
).send()
use_qdrant = False
use_qdrant_type = "Local"
if res and res.get("value") == "yes":
use_qdrant = True
local_res = await cl.AskActionMessage(
content="Do you want to use local or cloud?",
actions=[
cl.Action(name="Local", value="Local", label="✅ Local"),
cl.Action(name="Cloud", value="Cloud", label="❌ Cloud"),
],
).send()
if local_res and local_res.get("value") == "Cloud":
use_qdrant_type = "Cloud"
use_rct = False
res = await cl.AskActionMessage(
content="Do you want to use RecursiveCharacterTextSplitter?",
actions=[
cl.Action(name="yes", value="yes", label="✅ Yes"),
cl.Action(name="no", value="no", label="❌ No"),
],
).send()
if res and res.get("value") == "yes":
use_rct = True
files = None
# Wait for the user to upload a file
while not files:
files = await cl.AskFileMessage(
content="Please upload a .txt or .pdf file to begin processing!",
accept=["text/plain", "application/pdf"],
max_size_mb=2,
timeout=180,
).send()
file = files[0]
msg = cl.Message(
content=f"Processing `{file.name}`...", disable_human_feedback=True
)
await msg.send()
texts = process_file(file, use_rct)
msg = cl.Message(
content=f"Resulted in {len(texts)} chunks", disable_human_feedback=True
)
await msg.send()
# decide if to use the dict vector store of the Qdrant vector store
# Create a dict vector store
if use_qdrant == False:
vector_db = VectorDatabase()
vector_db = await vector_db.abuild_from_list(texts)
else:
embedding_model = EmbeddingModel(embeddings_model_name= "text-embedding-3-small", dimensions=1000)
if use_qdrant_type == "Local":
from utilities.vector_database import QdrantDatabase
vector_db = QdrantDatabase(
embedding_model=embedding_model
)
vector_db = await vector_db.abuild_from_list(texts)
msg = cl.Message(
content=f"The Vector store has been created", disable_human_feedback=True
)
await msg.send()
chat_openai = ChatOpenAI()
# Create a chain
retrieval_augmented_qa_pipeline = RetrievalAugmentedQAPipeline(
vector_db_retriever=vector_db,
llm=chat_openai,
system_role_prompt=system_role_prompt,
user_role_prompt=user_role_prompt
)
# Let the user know that the system is ready
msg.content = f"Processing `{file.name}` is complete."
await msg.update()
msg.content = f"You can now ask questions about `{file.name}`."
await msg.update()
cl.user_session.set("chain", retrieval_augmented_qa_pipeline)
@cl.on_message
async def main(message):
chain = cl.user_session.get("chain")
msg = cl.Message(content="")
result = await chain.arun_pipeline(message.content)
async for stream_resp in result["response"]:
await msg.stream_token(stream_resp)
await msg.send() |